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

Filtrable hnsw (#26)

Browse Source 
* raw points scorer

* raw point scorer for memmap storage

* search interface prepare

* graph binary saving + store PointOffsetId as u32

* WIP: entry points

* connect new link method

* update libs + search layer method + visited list + search context + update rust

* implement Euclid metric + always use MinHeap for priority queue

* small refactor

* search for 0 level entry

* update visited pool to be lock free and thread safe

* use ef_construct from graph layer struct + limit visited links to M

* add metric pre-processing before on vector upsert

* old hnsw heuristic

* save hnsw graph for export

* search method + tests

* small fixes

* add benchmark and profiler

* build time optimizations

* use SeaHash

* remove unsed benchmark

* merge hnsw graph function

* WIP:HNSW index build function

* HNSW build_index with additional indexing

* refactor fixtures

* graph save and load test

* test and fixes for filterable HNSW

* enable hnsw index for query planning

* fix cardinality estimation tests + remove query planner as class

* small refactor

* store full copy of collection settings with collection + allow partial override on creation #16

* API for updating collection parameters #16

* refactor: move collection error -> types

* report collection status in info API #17

* update OpenAPI Schema
This commit is contained in:
Andrey Vasnetsov
2021-05-30 17:14:42 +02:00
committed by GitHub
parent c83ddec2cb
commit 3616631300
85 changed files with 5247 additions and 1822 deletions
Download Patch File Download Diff File Expand all files Collapse all files

924
Cargo.lock generated
View File

File diff suppressed because it is too large Load Diff

2
Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "qdrant"
version = "0.2.1"
version = "0.3.0"
authors = ["Andrey Vasnetsov <andrey@vasnetsov.com>"]
edition = "2018"
doctest = false

4
Dockerfile
View File

@@ -1,10 +1,10 @@
FROM rust:1.49 as builder
FROM rust:1.51 as builder
COPY . ./qdrant
WORKDIR ./qdrant
ENV OPENBLAS_TARGET=CORE2
RUN apt-get update ; apt-get install -y clang libopenblas-dev libgfortran-8-dev
RUN apt-get update ; apt-get install -y clang libopenblas-dev libgfortran-8-dev gfortran
# Build actual target here
RUN cargo build --release --bin qdrant

28
QUICK_START.md
View File

@@ -36,19 +36,33 @@ Expected response:
```json
{
"result": {
"status": "green",
"vectors_count": 0,
"segments_count": 5,
"disk_data_size": 0,
"ram_data_size": 0,
"config": {
"vector_size": 4,
"index": {
"type": "plain",
"options": {}
"params": {
"vector_size": 4,
"distance": "Dot"
},
"distance": "Dot",
"storage_type": {
"type": "in_memory"
"hnsw_config": {
"m": 16,
"ef_construct": 100,
"full_scan_threshold": 10000
},
"optimizer_config": {
"deleted_threshold": 0.2,
"vacuum_min_vector_number": 1000,
"max_segment_number": 5,
"memmap_threshold": 50000,
"indexing_threshold": 20000,
"payload_indexing_threshold": 10000,
"flush_interval_sec": 1
},
"wal_config": {
"wal_capacity_mb": 32,
"wal_segments_ahead": 0
}
}
},

10
config/config.yaml
View File

@@ -39,6 +39,16 @@ storage:
# Minimum interval between forced flushes.
flush_interval_sec: 10
# Default parameters of HNSW Index. Could be override for each collection individually
hnsw_index:
# Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.
m: 16
# Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.
ef_construct: 100
# Minimal amount of points for additional payload-based indexing.
# If payload chunk is smaller than `full_scan_threshold` additional indexing won't be used -
# in this case full-scan search should be preferred by query planner and additional indexing is not required.
full_scan_threshold: 10000
service:

544
docs/redoc/openapi.json
View File

@@ -4,6 +4,7 @@
"AliasOperations": {
"anyOf": [
{
"additionalProperties": false,
"description": "Create alternative name for a collection. Collection will be available under both names for search, retrieve,",
"properties": {
"create_alias": {
@@ -28,6 +29,7 @@
"type": "object"
},
{
"additionalProperties": false,
"description": "Delete alias if exists",
"properties": {
"delete_alias": {
@@ -48,6 +50,7 @@
"type": "object"
},
{
"additionalProperties": false,
"description": "Change alias to a new one",
"properties": {
"rename_alias": {
@@ -73,6 +76,29 @@
}
]
},
"CollectionConfig": {
"properties": {
"hnsw_config": {
"$ref": "#/components/schemas/HnswConfig"
},
"optimizer_config": {
"$ref": "#/components/schemas/OptimizersConfig"
},
"params": {
"$ref": "#/components/schemas/CollectionParams"
},
"wal_config": {
"$ref": "#/components/schemas/WalConfig"
}
},
"required": [
"hnsw_config",
"optimizer_config",
"params",
"wal_config"
],
"type": "object"
},
"CollectionDescription": {
"properties": {
"name": {
@@ -88,7 +114,7 @@
"description": "Current statistics and configuration of the collection.",
"properties": {
"config": {
"$ref": "#/components/schemas/SegmentConfig"
"$ref": "#/components/schemas/CollectionConfig"
},
"disk_data_size": {
"description": "Disk space, used by collection",
@@ -108,6 +134,9 @@
"minimum": 0,
"type": "integer"
},
"status": {
"$ref": "#/components/schemas/CollectionStatus"
},
"vectors_count": {
"description": "Number of vectors in collection",
"format": "uint",
@@ -120,10 +149,37 @@
"disk_data_size",
"ram_data_size",
"segments_count",
"status",
"vectors_count"
],
"type": "object"
},
"CollectionParams": {
"properties": {
"distance": {
"$ref": "#/components/schemas/Distance"
},
"vector_size": {
"description": "Size of a vectors used",
"format": "uint",
"minimum": 0,
"type": "integer"
}
},
"required": [
"distance",
"vector_size"
],
"type": "object"
},
"CollectionStatus": {
"enum": [
"green",
"yellow",
"red"
],
"type": "string"
},
"CollectionUpdateOperations": {
"anyOf": [
{
@@ -254,6 +310,7 @@
"FieldIndexOperations": {
"anyOf": [
{
"additionalProperties": false,
"description": "Create index for payload field",
"properties": {
"create_index": {
@@ -266,6 +323,7 @@
"type": "object"
},
{
"additionalProperties": false,
"description": "Delete index for the field",
"properties": {
"delete_index": {
@@ -366,7 +424,8 @@
"minimum": 0,
"type": "integer"
},
"type": "array"
"type": "array",
"uniqueItems": true
}
},
"required": [
@@ -374,65 +433,59 @@
],
"type": "object"
},
"Indexes": {
"anyOf": [
{
"description": "Do not use any index, scan whole vector collection during search. Guarantee 100% precision, but may be time consuming on large collections.",
"properties": {
"options": {
"type": "object"
},
"type": {
"enum": [
"plain"
],
"type": "string"
}
},
"required": [
"options",
"type"
],
"type": "object"
"HnswConfig": {
"properties": {
"ef_construct": {
"description": "Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.",
"format": "uint",
"minimum": 0,
"type": "integer"
},
{
"description": "Use filterable HNSW index for approximate search. Is very fast even on a very huge collections, but require additional space to store index and additional time to build it.",
"properties": {
"options": {
"properties": {
"ef_construct": {
"description": "Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.",
"format": "uint",
"minimum": 0,
"type": "integer"
},
"m": {
"description": "Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.",
"format": "uint",
"minimum": 0,
"type": "integer"
}
},
"required": [
"ef_construct",
"m"
],
"type": "object"
},
"type": {
"enum": [
"hnsw"
],
"type": "string"
}
},
"required": [
"options",
"type"
],
"type": "object"
"full_scan_threshold": {
"description": "Minimal amount of points for additional payload-based indexing. If payload chunk is smaller than `full_scan_threshold` additional indexing won't be used - in this case full-scan search should be preferred by query planner and additional indexing is not required.",
"format": "uint",
"minimum": 0,
"type": "integer"
},
"m": {
"description": "Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.",
"format": "uint",
"minimum": 0,
"type": "integer"
}
]
},
"required": [
"ef_construct",
"full_scan_threshold",
"m"
],
"type": "object"
},
"HnswConfigDiff": {
"properties": {
"ef_construct": {
"description": "Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
},
"full_scan_threshold": {
"description": "Minimal amount of points for additional payload-based indexing. If payload chunk is smaller than `full_scan_threshold` additional indexing won't be used - in this case full-scan search should be preferred by query planner and additional indexing is not required.",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
},
"m": {
"description": "Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
}
},
"type": "object"
},
"Match": {
"properties": {
@@ -450,42 +503,131 @@
},
"type": "object"
},
"PayloadIndexType": {
"anyOf": [
{
"description": "Do not index anything, just keep of what should be indexed later",
"properties": {
"type": {
"enum": [
"plain"
],
"type": "string"
}
},
"required": [
"type"
],
"type": "object"
"OptimizersConfig": {
"properties": {
"deleted_threshold": {
"description": "The minimal fraction of deleted vectors in a segment, required to perform segment optimization",
"format": "double",
"type": "number"
},
{
"description": "Build payload index. Index is saved on disc, but index itself is in RAM",
"properties": {
"type": {
"enum": [
"struct"
],
"type": "string"
}
},
"required": [
"type"
],
"type": "object"
"flush_interval_sec": {
"description": "Minimum interval between forced flushes.",
"format": "uint64",
"minimum": 0,
"type": "integer"
},
"indexing_threshold": {
"description": "Maximum number of vectors allowed for plain index. Default value based on https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md",
"format": "uint",
"minimum": 0,
"type": "integer"
},
"max_segment_number": {
"description": "If the number of segments exceeds this value, the optimizer will merge the smallest segments.",
"format": "uint",
"minimum": 0,
"type": "integer"
},
"memmap_threshold": {
"description": "Maximum number of vectors to store in-memory per segment. Segments larger than this threshold will be stored as read-only memmaped file.",
"format": "uint",
"minimum": 0,
"type": "integer"
},
"payload_indexing_threshold": {
"description": "Starting from this amount of vectors per-segment the engine will start building index for payload.",
"format": "uint",
"minimum": 0,
"type": "integer"
},
"vacuum_min_vector_number": {
"description": "The minimal number of vectors in a segment, required to perform segment optimization",
"format": "uint",
"minimum": 0,
"type": "integer"
}
},
"required": [
"deleted_threshold",
"flush_interval_sec",
"indexing_threshold",
"max_segment_number",
"memmap_threshold",
"payload_indexing_threshold",
"vacuum_min_vector_number"
],
"description": "Type of payload index"
"type": "object"
},
"OptimizersConfigDiff": {
"properties": {
"deleted_threshold": {
"description": "The minimal fraction of deleted vectors in a segment, required to perform segment optimization",
"format": "double",
"nullable": true,
"type": "number"
},
"flush_interval_sec": {
"description": "Minimum interval between forced flushes.",
"format": "uint64",
"minimum": 0,
"nullable": true,
"type": "integer"
},
"indexing_threshold": {
"description": "Maximum number of vectors allowed for plain index. Default value based on https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
},
"max_segment_number": {
"description": "If the number of segments exceeds this value, the optimizer will merge the smallest segments.",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
},
"memmap_threshold": {
"description": "Maximum number of vectors to store in-memory per segment. Segments larger than this threshold will be stored as read-only memmaped file.",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
},
"payload_indexing_threshold": {
"description": "Starting from this amount of vectors per-segment the engine will start building index for payload.",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
},
"vacuum_min_vector_number": {
"description": "The minimal number of vectors in a segment, required to perform segment optimization",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
}
},
"type": "object"
},
"PayloadInterface": {
"anyOf": [
{
"$ref": "#/components/schemas/PayloadVariant_for_String"
},
{
"$ref": "#/components/schemas/PayloadVariant_for_int64"
},
{
"$ref": "#/components/schemas/PayloadVariant_for_double"
},
{
"$ref": "#/components/schemas/PayloadInterfaceStrict"
}
]
},
"PayloadInterfaceStrict": {
"anyOf": [
{
"properties": {
@@ -564,6 +706,7 @@
"PayloadOps": {
"anyOf": [
{
"additionalProperties": false,
"description": "Set payload value, overrides if it is already exists",
"properties": {
"set_payload": {
@@ -597,6 +740,7 @@
"type": "object"
},
{
"additionalProperties": false,
"description": "Deletes specified payload values if they are assigned",
"properties": {
"delete_payload": {
@@ -630,6 +774,7 @@
"type": "object"
},
{
"additionalProperties": false,
"description": "Drops all Payload values associated with given points.",
"properties": {
"clear_payload": {
@@ -806,6 +951,7 @@
"PointInsertOperations": {
"anyOf": [
{
"additionalProperties": false,
"description": "Inset points from a batch.",
"properties": {
"batch": {
@@ -853,6 +999,7 @@
"type": "object"
},
{
"additionalProperties": false,
"description": "Insert points from a list",
"properties": {
"points": {
@@ -872,6 +1019,7 @@
"PointOperations": {
"anyOf": [
{
"additionalProperties": false,
"description": "Insert or update points",
"properties": {
"upsert_points": {
@@ -884,6 +1032,7 @@
"type": "object"
},
{
"additionalProperties": false,
"description": "Delete point if exists",
"properties": {
"delete_points": {
@@ -1096,32 +1245,17 @@
"type": "object"
},
"SearchParams": {
"anyOf": [
{
"description": "Params relevant to HNSW index",
"properties": {
"hnsw": {
"properties": {
"ef": {
"description": "Size of the beam in a beam-search. Larger the value - more accurate the result, more time required for search.",
"format": "uint",
"minimum": 0,
"type": "integer"
}
},
"required": [
"ef"
],
"type": "object"
}
},
"required": [
"hnsw"
],
"type": "object"
"description": "Additional parameters of the search",
"properties": {
"hnsw_ef": {
"description": "Params relevant to HNSW index /// Size of the beam in a beam-search. Larger the value - more accurate the result, more time required for search.",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
}
],
"description": "Additional parameters of the search"
},
"type": "object"
},
"SearchRequest": {
"description": "Search request",
@@ -1169,46 +1303,10 @@
],
"type": "object"
},
"SegmentConfig": {
"properties": {
"distance": {
"$ref": "#/components/schemas/Distance"
},
"index": {
"$ref": "#/components/schemas/Indexes"
},
"payload_index": {
"anyOf": [
{
"$ref": "#/components/schemas/PayloadIndexType"
},
{
"nullable": true
}
],
"description": "Payload Indexes"
},
"storage_type": {
"$ref": "#/components/schemas/StorageType"
},
"vector_size": {
"description": "Size of a vectors used",
"format": "uint",
"minimum": 0,
"type": "integer"
}
},
"required": [
"distance",
"index",
"storage_type",
"vector_size"
],
"type": "object"
},
"StorageOperations": {
"anyOf": [
{
"additionalProperties": false,
"description": "Create new collection and (optionally) specify index params",
"properties": {
"create_collection": {
@@ -1216,23 +1314,46 @@
"distance": {
"$ref": "#/components/schemas/Distance"
},
"index": {
"hnsw_config": {
"anyOf": [
{
"$ref": "#/components/schemas/Indexes"
"$ref": "#/components/schemas/HnswConfigDiff"
},
{
"nullable": true
}
]
],
"description": "Custom params for HNSW index. If none - values from service configuration file are used."
},
"name": {
"type": "string"
},
"optimizers_config": {
"anyOf": [
{
"$ref": "#/components/schemas/OptimizersConfigDiff"
},
{
"nullable": true
}
],
"description": "Custom params for Optimizers. If none - values from service configuration file are used."
},
"vector_size": {
"format": "uint",
"minimum": 0,
"type": "integer"
},
"wal_config": {
"anyOf": [
{
"$ref": "#/components/schemas/WalConfigDiff"
},
{
"nullable": true
}
],
"description": "Custom params for WAL. If none - values from service configuration file are used."
}
},
"required": [
@@ -1249,6 +1370,39 @@
"type": "object"
},
{
"additionalProperties": false,
"description": "Update parameters of the existing collection",
"properties": {
"update_collection": {
"properties": {
"name": {
"type": "string"
},
"optimizers_config": {
"anyOf": [
{
"$ref": "#/components/schemas/OptimizersConfigDiff"
},
{
"nullable": true
}
],
"description": "Custom params for Optimizers. If none - values from service configuration file are used. This operation is blocking, it will only proceed ones all current optimizations are complete"
}
},
"required": [
"name"
],
"type": "object"
}
},
"required": [
"update_collection"
],
"type": "object"
},
{
"additionalProperties": false,
"description": "Delete collection with given name",
"properties": {
"delete_collection": {
@@ -1261,6 +1415,7 @@
"type": "object"
},
{
"additionalProperties": false,
"description": "Perform changes of collection aliases. Alias changes are atomic, meaning that no collection modifications can happen between alias operations.",
"properties": {
"change_aliases": {
@@ -1285,41 +1440,6 @@
}
]
},
"StorageType": {
"anyOf": [
{
"description": "Store vectors in memory and use persistence storage only if vectors are changed",
"properties": {
"type": {
"enum": [
"in_memory"
],
"type": "string"
}
},
"required": [
"type"
],
"type": "object"
},
{
"description": "Use memmap to store vectors, a little slower than `InMemory`, but requires little RAM",
"properties": {
"type": {
"enum": [
"mmap"
],
"type": "string"
}
},
"required": [
"type"
],
"type": "object"
}
],
"description": "Type of vector storage"
},
"UpdateResult": {
"properties": {
"operation_id": {
@@ -1344,6 +1464,46 @@
"completed"
],
"type": "string"
},
"WalConfig": {
"properties": {
"wal_capacity_mb": {
"description": "Size of a single WAL segment in MB",
"format": "uint",
"minimum": 0,
"type": "integer"
},
"wal_segments_ahead": {
"description": "Number of WAL segments to create ahead of actually used ones",
"format": "uint",
"minimum": 0,
"type": "integer"
}
},
"required": [
"wal_capacity_mb",
"wal_segments_ahead"
],
"type": "object"
},
"WalConfigDiff": {
"properties": {
"wal_capacity_mb": {
"description": "Size of a single WAL segment in MB",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
},
"wal_segments_ahead": {
"description": "Number of WAL segments to create ahead of actually used ones",
"format": "uint",
"minimum": 0,
"nullable": true,
"type": "integer"
}
},
"type": "object"
}
}
},
@@ -1355,13 +1515,13 @@
"contact": {
"email": "andrey@vasnetsov.com"
},
"description": "\nAPI description for Qdrant vector search engine.\n\nThis document describes CRUD and search operations on collections of points (vectors with payload).\n\nQdrant supports any combinations of `should`, `must` and `must_not` conditions, which makes it possible to use in applications when object could not be described solely by vector. It could be location features, availability flags, and other custom properties businesses should take into account.\n## Examples\nThis examples cover the most basic use-cases - collection creation and basic vector search.\n### Create collection\nFirst - let's create a collection with dot-production metric.\n```\ncurl -X POST 'http://localhost:6333/collections' \\\n -H 'Content-Type: application/json' \\\n --data-raw '{\n \"create_collection\": {\n \"name\": \"test_collection\",\n \"vector_size\": 4,\n \"distance\": \"Dot\"\n }\n }'\n```\nExpected response:\n```\n{\n \"result\": true,\n \"status\": \"ok\",\n \"time\": 0.031095451\n}\n```\nWe can ensure that collection was created:\n```\ncurl 'http://localhost:6333/collections/test_collection'\n```\nExpected response:\n```\n{\n \"result\": {\n \"vectors_count\": 0,\n \"segments_count\": 5,\n \"disk_data_size\": 0,\n \"ram_data_size\": 0,\n \"config\": {\n \"vector_size\": 4,\n \"index\": {\n \"type\": \"plain\",\n \"options\": {}\n },\n \"distance\": \"Dot\",\n \"storage_type\": {\n \"type\": \"in_memory\"\n }\n }\n },\n \"status\": \"ok\",\n \"time\": 2.1199e-05\n}\n```\n\n### Add points\nLet's now add vectors with some payload:\n```\ncurl -L -X POST 'http://localhost:6333/collections/test_collection?wait=true' \\ -H 'Content-Type: application/json' \\ --data-raw '{\n \"upsert_points\": {\n \"points\": [\n {\"id\": 1, \"vector\": [0.05, 0.61, 0.76, 0.74], \"payload\": {\"city\": {\"type\": \"keyword\", \"value\": \"Berlin\"}}},\n {\"id\": 2, \"vector\": [0.19, 0.81, 0.75, 0.11], \"payload\": {\"city\": {\"type\": \"keyword\", \"value\": [\"Berlin\", \"London\"] }}},\n {\"id\": 3, \"vector\": [0.36, 0.55, 0.47, 0.94], \"payload\": {\"city\": {\"type\": \"keyword\", \"value\": [\"Berlin\", \"Moscow\"] }}},\n {\"id\": 4, \"vector\": [0.18, 0.01, 0.85, 0.80], \"payload\": {\"city\": {\"type\": \"keyword\", \"value\": [\"London\", \"Moscow\"]}}},\n {\"id\": 5, \"vector\": [0.24, 0.18, 0.22, 0.44], \"payload\": {\"count\": {\"type\": \"integer\", \"value\": [0]}}},\n {\"id\": 6, \"vector\": [0.35, 0.08, 0.11, 0.44]}\n ]\n }\n}'\n```\nExpected response:\n```\n{\n \"result\": {\n \"operation_id\": 0,\n \"status\": \"completed\"\n },\n \"status\": \"ok\",\n \"time\": 0.000206061\n}\n```\n### Search with filtering\nLet's start with a basic request:\n```\ncurl -L -X POST 'http://localhost:6333/collections/test_collection/points/search' \\ -H 'Content-Type: application/json' \\ --data-raw '{\n \"vector\": [0.2,0.1,0.9,0.7],\n \"top\": 3\n}'\n```\nExpected response:\n```\n{\n \"result\": [\n { \"id\": 4, \"score\": 1.362 },\n { \"id\": 1, \"score\": 1.273 },\n { \"id\": 3, \"score\": 1.208 }\n ],\n \"status\": \"ok\",\n \"time\": 0.000055785\n}\n```\nBut result is different if we add a filter:\n```\ncurl -L -X POST 'http://localhost:6333/collections/test_collection/points/search' \\ -H 'Content-Type: application/json' \\ --data-raw '{\n \"filter\": {\n \"should\": [\n {\n \"key\": \"city\",\n \"match\": {\n \"keyword\": \"London\"\n }\n }\n ]\n },\n \"vector\": [0.2, 0.1, 0.9, 0.7],\n \"top\": 3\n}'\n```\nExpected response:\n```\n{\n \"result\": [\n { \"id\": 4, \"score\": 1.362 },\n { \"id\": 2, \"score\": 0.871 }\n ],\n \"status\": \"ok\",\n \"time\": 0.000093972\n}\n```\n",
"description": "\nAPI description for Qdrant vector search engine.\n\nThis document describes CRUD and search operations on collections of points (vectors with payload).\n\nQdrant supports any combinations of `should`, `must` and `must_not` conditions, which makes it possible to use in applications when object could not be described solely by vector. It could be location features, availability flags, and other custom properties businesses should take into account.\n## Examples\nThis examples cover the most basic use-cases - collection creation and basic vector search.\n### Create collection\nFirst - let's create a collection with dot-production metric.\n```\ncurl -X POST 'http://localhost:6333/collections' \\\n -H 'Content-Type: application/json' \\\n --data-raw '{\n \"create_collection\": {\n \"name\": \"test_collection\",\n \"vector_size\": 4,\n \"distance\": \"Dot\"\n }\n }'\n```\nExpected response:\n```\n{\n \"result\": true,\n \"status\": \"ok\",\n \"time\": 0.031095451\n}\n```\nWe can ensure that collection was created:\n```\ncurl 'http://localhost:6333/collections/test_collection'\n```\nExpected response:\n```\n{\n \"result\": {\n \"status\": \"green\",\n \"vectors_count\": 0,\n \"segments_count\": 5,\n \"disk_data_size\": 0,\n \"ram_data_size\": 0,\n \"config\": {\n \"params\": {\n \"vector_size\": 4,\n \"distance\": \"Dot\"\n },\n \"hnsw_config\": {\n \"m\": 16,\n \"ef_construct\": 100,\n \"full_scan_threshold\": 10000\n },\n \"optimizer_config\": {\n \"deleted_threshold\": 0.2,\n \"vacuum_min_vector_number\": 1000,\n \"max_segment_number\": 5,\n \"memmap_threshold\": 50000,\n \"indexing_threshold\": 20000,\n \"payload_indexing_threshold\": 10000,\n \"flush_interval_sec\": 1\n },\n \"wal_config\": {\n \"wal_capacity_mb\": 32,\n \"wal_segments_ahead\": 0\n }\n }\n },\n \"status\": \"ok\",\n \"time\": 2.1199e-05\n}\n```\n\n### Add points\nLet's now add vectors with some payload:\n```\ncurl -L -X POST 'http://localhost:6333/collections/test_collection?wait=true' \\ -H 'Content-Type: application/json' \\ --data-raw '{\n \"upsert_points\": {\n \"points\": [\n {\"id\": 1, \"vector\": [0.05, 0.61, 0.76, 0.74], \"payload\": {\"city\": {\"type\": \"keyword\", \"value\": \"Berlin\"}}},\n {\"id\": 2, \"vector\": [0.19, 0.81, 0.75, 0.11], \"payload\": {\"city\": {\"type\": \"keyword\", \"value\": [\"Berlin\", \"London\"] }}},\n {\"id\": 3, \"vector\": [0.36, 0.55, 0.47, 0.94], \"payload\": {\"city\": {\"type\": \"keyword\", \"value\": [\"Berlin\", \"Moscow\"] }}},\n {\"id\": 4, \"vector\": [0.18, 0.01, 0.85, 0.80], \"payload\": {\"city\": {\"type\": \"keyword\", \"value\": [\"London\", \"Moscow\"]}}},\n {\"id\": 5, \"vector\": [0.24, 0.18, 0.22, 0.44], \"payload\": {\"count\": {\"type\": \"integer\", \"value\": [0]}}},\n {\"id\": 6, \"vector\": [0.35, 0.08, 0.11, 0.44]}\n ]\n }\n}'\n```\nExpected response:\n```\n{\n \"result\": {\n \"operation_id\": 0,\n \"status\": \"completed\"\n },\n \"status\": \"ok\",\n \"time\": 0.000206061\n}\n```\n### Search with filtering\nLet's start with a basic request:\n```\ncurl -L -X POST 'http://localhost:6333/collections/test_collection/points/search' \\ -H 'Content-Type: application/json' \\ --data-raw '{\n \"vector\": [0.2,0.1,0.9,0.7],\n \"top\": 3\n}'\n```\nExpected response:\n```\n{\n \"result\": [\n { \"id\": 4, \"score\": 1.362 },\n { \"id\": 1, \"score\": 1.273 },\n { \"id\": 3, \"score\": 1.208 }\n ],\n \"status\": \"ok\",\n \"time\": 0.000055785\n}\n```\nBut result is different if we add a filter:\n```\ncurl -L -X POST 'http://localhost:6333/collections/test_collection/points/search' \\ -H 'Content-Type: application/json' \\ --data-raw '{\n \"filter\": {\n \"should\": [\n {\n \"key\": \"city\",\n \"match\": {\n \"keyword\": \"London\"\n }\n }\n ]\n },\n \"vector\": [0.2, 0.1, 0.9, 0.7],\n \"top\": 3\n}'\n```\nExpected response:\n```\n{\n \"result\": [\n { \"id\": 4, \"score\": 1.362 },\n { \"id\": 2, \"score\": 0.871 }\n ],\n \"status\": \"ok\",\n \"time\": 0.000093972\n}\n```\n",
"license": {
"name": "Apache 2.0",
"url": "http://www.apache.org/licenses/LICENSE-2.0.html"
},
"title": "Qdrant API",
"version": "0.2.1"
"version": "0.3.0"
},
"openapi": "3.0.1",
"paths": {

3
lib/collection/Cargo.toml
View File

@@ -1,6 +1,6 @@
[package]
name = "collection"
version = "0.1.0"
version = "0.3.0"
authors = ["Andrey Vasnetsov <vasnetsov93@gmail.com>"]
edition = "2018"
@@ -30,6 +30,7 @@ crossbeam-channel = "0.4.3"
atomicwrites = "0.2.5"
log = "0.4"
env_logger = "0.7.1"
merge = "0.1.0"
segment = {path = "../segment"}

120
lib/collection/src/collection.rs
View File

@@ -1,77 +1,36 @@
use thiserror::Error;
use crate::operations::CollectionUpdateOperations;
use segment::types::{PointIdType, ScoredPoint, SegmentConfig, VectorElementType, HasIdCondition};
use std::result;
use crate::operations::types::{Record, CollectionInfo, UpdateResult, UpdateStatus, SearchRequest, RecommendRequest};
use std::sync::Arc;
use crate::wal::{SerdeWal, WalError};
use crate::segment_manager::segment_managers::{SegmentSearcher, SegmentUpdater};
use segment::entry::entry_point::OperationError;
use tokio::task::JoinError;
use crossbeam_channel::{Sender, SendError};
use crate::update_handler::update_handler::{UpdateHandler, UpdateSignal};
use parking_lot::{Mutex, RwLock};
use crate::segment_manager::holders::segment_holder::SegmentHolder;
use tokio::runtime::Runtime;
use itertools::Itertools;
use std::collections::HashMap;
use segment::types::Filter;
use std::path::PathBuf;
use std::sync::Arc;
use crossbeam_channel::{Sender};
use itertools::Itertools;
use parking_lot::{Mutex, RwLock};
use tokio::runtime::Runtime;
use segment::types::{HasIdCondition, PointIdType, ScoredPoint, VectorElementType, SegmentType};
use segment::types::Condition;
use segment::types::Filter;
#[derive(Error, Debug, Clone)]
#[error("{0}")]
pub enum CollectionError {
#[error("Wrong input: {description}")]
BadInput { description: String },
#[error("No point with id {missed_point_id} found")]
NotFound { missed_point_id: PointIdType },
#[error("Service internal error: {error}")]
ServiceError { error: String },
#[error("Bad request: {description}")]
BadRequest { description: String },
}
impl From<OperationError> for CollectionError {
fn from(err: OperationError) -> Self {
match err {
OperationError::WrongVector { .. } => Self::BadInput { description: format!("{}", err) },
OperationError::PointIdError { missed_point_id } => Self::NotFound { missed_point_id },
OperationError::ServiceError { description } => Self::ServiceError { error: description },
OperationError::TypeError { .. } => Self::BadInput { description: format!("{}", err) },
}
}
}
impl From<JoinError> for CollectionError {
fn from(err: JoinError) -> Self {
Self::ServiceError { error: format!("{}", err) }
}
}
impl From<WalError> for CollectionError {
fn from(err: WalError) -> Self {
Self::ServiceError { error: format!("{}", err) }
}
}
impl<T> From<SendError<T>> for CollectionError {
fn from(_err: SendError<T>) -> Self {
Self::ServiceError { error: format!("Can't reach one of the workers") }
}
}
pub type CollectionResult<T> = result::Result<T, CollectionError>;
use crate::collection_builder::optimizers_builder::build_optimizers;
use crate::config::CollectionConfig;
use crate::operations::CollectionUpdateOperations;
use crate::operations::config_diff::{DiffConfig, OptimizersConfigDiff};
use crate::operations::types::{CollectionError, CollectionInfo, CollectionResult, RecommendRequest, Record, SearchRequest, UpdateResult, UpdateStatus, CollectionStatus};
use crate::segment_manager::holders::segment_holder::SegmentHolder;
use crate::segment_manager::segment_managers::{SegmentSearcher, SegmentUpdater};
use crate::update_handler::update_handler::{UpdateHandler, UpdateSignal};
use crate::wal::{SerdeWal};
pub struct Collection {
pub segments: Arc<RwLock<SegmentHolder>>,
pub config: SegmentConfig,
pub config: Arc<RwLock<CollectionConfig>>,
pub wal: Arc<Mutex<SerdeWal<CollectionUpdateOperations>>>,
pub searcher: Arc<dyn SegmentSearcher + Sync + Send>,
pub update_handler: Arc<UpdateHandler>,
pub update_handler: Arc<Mutex<UpdateHandler>>,
pub updater: Arc<dyn SegmentUpdater + Sync + Send>,
pub runtime_handle: Arc<Runtime>,
pub update_sender: Sender<UpdateSignal>,
pub path: PathBuf
}
@@ -107,19 +66,24 @@ impl Collection {
let mut segments_count = 0;
let mut ram_size = 0;
let mut disk_size = 0;
let mut status = CollectionStatus::Green;
for (_idx, segment) in segments.iter() {
segments_count += 1;
let segment_info = segment.get().read().info();
if segment_info.segment_type == SegmentType::Special {
status = CollectionStatus::Yellow;
}
vectors_count += segment_info.num_vectors;
disk_size += segment_info.disk_usage_bytes;
ram_size += segment_info.ram_usage_bytes;
}
Ok(CollectionInfo {
status,
vectors_count,
segments_count,
disk_data_size: disk_size,
ram_data_size: ram_size,
config: self.config.clone(),
config: self.config.read().clone(),
})
}
@@ -230,6 +194,34 @@ impl Collection {
self.search(Arc::new(search_request))
}
/// Updates collection optimization params:
/// - Saves new params on disk
/// - Stops existing optimization loop
/// - Runs new optimizers with new params
pub fn update_optimizer_params(&self, optimizer_config_diff: OptimizersConfigDiff) -> CollectionResult<()> {
{
let mut config = self.config.write();
config.optimizer_config = optimizer_config_diff.update(&config.optimizer_config)?;
config.save(self.path.as_path())?;
}
let config = self.config.read();
let mut update_handler = self.update_handler.lock();
self.stop()?;
update_handler.wait_worker_stops()?;
let new_optimizers = build_optimizers(
self.path.as_path(),
&config.params,
&config.optimizer_config,
&config.hnsw_config
);
update_handler.optimizers = new_optimizers;
update_handler.flush_timeout_sec = config.optimizer_config.flush_interval_sec;
update_handler.run_worker();
self.update_sender.send(UpdateSignal::Nop)?;
Ok(())
}
}
impl Drop for Collection {

102
lib/collection/src/collection_builder/collection_builder.rs
View File

@@ -1,52 +1,34 @@
use crate::collection::{Collection, CollectionResult, CollectionError};
use crate::segment_manager::holders::segment_holder::SegmentHolder;
use segment::segment_constructor::simple_segment_constructor::build_simple_segment;
use std::path::Path;
use crate::wal::SerdeWal;
use crate::operations::CollectionUpdateOperations;
use wal::WalOptions;
use std::sync::Arc;
use tokio::runtime::Runtime;
use crate::segment_manager::simple_segment_searcher::SimpleSegmentSearcher;
use crate::segment_manager::simple_segment_updater::SimpleSegmentUpdater;
use crossbeam_channel::unbounded;
use crate::update_handler::update_handler::{UpdateHandler, Optimizer};
use segment::types::SegmentConfig;
use std::fs::create_dir_all;
use parking_lot::{RwLock, Mutex};
use std::path::Path;
use std::sync::Arc;
use crossbeam_channel::unbounded;
use parking_lot::{Mutex, RwLock};
use tokio::runtime;
use tokio::runtime::Runtime;
use segment::segment_constructor::simple_segment_constructor::build_simple_segment;
use segment::types::HnswConfig;
use crate::collection::Collection;
use crate::collection_builder::optimizers_builder::build_optimizers;
use crate::collection_builder::optimizers_builder::OptimizersConfig;
use atomicwrites::OverwriteBehavior::AllowOverwrite;
use atomicwrites::AtomicFile;
use std::io::Write;
use tokio::runtime;
const DEFAULT_SEGMENT_NUMBER: usize = 5;
pub const COLLECTION_CONFIG_FILE: &str = "config.json";
fn save_config(path: &Path, config: &SegmentConfig) -> CollectionResult<()> {
let config_path = path.join(COLLECTION_CONFIG_FILE);
let af = AtomicFile::new(&config_path, AllowOverwrite);
let state_bytes = serde_json::to_vec(config).unwrap();
af.write(|f| {
f.write_all(&state_bytes)
}).or_else(move |err|
Err(CollectionError::ServiceError {
error: format!("Can't write {:?}, error: {}", config_path, err)
})
)?;
Ok(())
}
use crate::config::{CollectionConfig, CollectionParams, WalConfig};
use crate::operations::CollectionUpdateOperations;
use crate::operations::types::{CollectionError, CollectionResult};
use crate::segment_manager::holders::segment_holder::SegmentHolder;
use crate::segment_manager::simple_segment_searcher::SimpleSegmentSearcher;
use crate::segment_manager::simple_segment_updater::SimpleSegmentUpdater;
use crate::update_handler::update_handler::{Optimizer, UpdateHandler};
use crate::wal::SerdeWal;
pub fn construct_collection(
segment_holder: SegmentHolder,
config: &SegmentConfig,
config: CollectionConfig,
wal: SerdeWal<CollectionUpdateOperations>,
search_runtime: Arc<Runtime>, // from service
optimizers: Arc<Vec<Box<Optimizer>>>,
flush_interval_sec: u64,
collection_path: &Path,
) -> Collection {
let segment_holder = Arc::new(RwLock::new(segment_holder));
@@ -62,27 +44,28 @@ pub fn construct_collection(
);
let updater = SimpleSegmentUpdater::new(segment_holder.clone());
// ToDo: Move tx-rx into updater, so Collection should not know about it.
let (tx, rx) = unbounded();
let update_handler = Arc::new(UpdateHandler::new(
let update_handler = Arc::new(Mutex::new(UpdateHandler::new(
optimizers,
rx,
optimize_runtime.clone(),
segment_holder.clone(),
locked_wal.clone(),
flush_interval_sec,
));
config.optimizer_config.flush_interval_sec,
)));
let collection = Collection {
segments: segment_holder.clone(),
config: config.clone(),
config: Arc::new(RwLock::new(config)),
wal: locked_wal,
searcher: Arc::new(searcher),
update_handler,
updater: Arc::new(updater),
runtime_handle: optimize_runtime,
update_sender: tx,
path: collection_path.to_owned()
};
return collection;
@@ -92,10 +75,11 @@ pub fn construct_collection(
/// Creates new empty collection with given configuration
pub fn build_collection(
collection_path: &Path,
wal_options: &WalOptions, // from config
segment_config: &SegmentConfig, // from user
wal_config: &WalConfig, // from config
collection_params: &CollectionParams, // from user
search_runtime: Arc<Runtime>, // from service
optimizers_config: &OptimizersConfig,
hnsw_config: &HnswConfig,
) -> CollectionResult<Collection> {
let wal_path = collection_path
.join("wal");
@@ -114,31 +98,39 @@ pub fn build_collection(
let mut segment_holder = SegmentHolder::new();
for _sid in 0..DEFAULT_SEGMENT_NUMBER {
for _sid in 0..optimizers_config.max_segment_number {
let segment = build_simple_segment(
segments_path.as_path(),
segment_config.vector_size,
segment_config.distance.clone())?;
collection_params.vector_size,
collection_params.distance)?;
segment_holder.add(segment);
}
let wal: SerdeWal<CollectionUpdateOperations> = SerdeWal::new(wal_path.to_str().unwrap(), wal_options)?;
let wal: SerdeWal<CollectionUpdateOperations> = SerdeWal::new(wal_path.to_str().unwrap(), &wal_config.into())?;
save_config(collection_path, &segment_config)?;
let collection_config = CollectionConfig {
params: collection_params.clone(),
hnsw_config: hnsw_config.clone(),
optimizer_config: optimizers_config.clone(),
wal_config: wal_config.clone()
};
collection_config.save(collection_path)?;
let optimizers = build_optimizers(
collection_path,
&segment_config,
&collection_params,
&optimizers_config,
&collection_config.hnsw_config
);
let collection = construct_collection(
segment_holder,
segment_config,
collection_config,
wal,
search_runtime,
optimizers,
optimizers_config.flush_interval_sec,
collection_path,
);
Ok(collection)

59
lib/collection/src/collection_builder/collection_loader.rs
View File

@@ -1,41 +1,36 @@
use crate::collection::{Collection, CollectionError};
use std::fs::read_dir;
use std::path::Path;
use tokio::runtime::Runtime;
use crate::segment_manager::holders::segment_holder::SegmentHolder;
use crate::wal::SerdeWal;
use crate::operations::CollectionUpdateOperations;
use wal::WalOptions;
use std::fs::{read_dir, File};
use segment::segment_constructor::segment_constructor::load_segment;
use crate::collection_builder::collection_builder::{construct_collection, COLLECTION_CONFIG_FILE};
use indicatif::ProgressBar;
use crate::collection_builder::optimizers_builder::OptimizersConfig;
use crate::collection_builder::optimizers_builder::build_optimizers;
use segment::types::SegmentConfig;
use std::io::Read;
use std::sync::Arc;
use indicatif::ProgressBar;
use tokio::runtime::Runtime;
fn load_config(path: &Path) -> SegmentConfig {
let config_path = path.join(COLLECTION_CONFIG_FILE);
let mut contents = String::new();
let mut file = File::open(config_path).unwrap();
file.read_to_string(&mut contents).unwrap();
serde_json::from_str(&contents).unwrap()
}
use segment::segment_constructor::segment_constructor::load_segment;
use crate::collection::Collection;
use crate::collection_builder::collection_builder::construct_collection;
use crate::collection_builder::optimizers_builder::build_optimizers;
use crate::config::CollectionConfig;
use crate::operations::CollectionUpdateOperations;
use crate::operations::types::CollectionError;
use crate::segment_manager::holders::segment_holder::SegmentHolder;
use crate::wal::SerdeWal;
pub fn load_collection(
collection_path: &Path,
wal_options: &WalOptions, // from config
search_runtime: Arc<Runtime>, // from service
optimizers_config: &OptimizersConfig,
) -> Collection {
let wal_path = collection_path.join("wal");
let segments_path = collection_path.join("segments");
let mut segment_holder = SegmentHolder::new();
let wal: SerdeWal<CollectionUpdateOperations> = SerdeWal::new(wal_path.to_str().unwrap(), wal_options).expect("Can't read WAL");
let collection_config = CollectionConfig::load(&collection_path)
.expect(&format!("Can't read collection config at {}", collection_path.to_str().unwrap()));
let wal: SerdeWal<CollectionUpdateOperations> = SerdeWal::new(
wal_path.to_str().unwrap(),
&(&collection_config.wal_config).into()
).expect("Can't read WAL");
let segment_dirs = read_dir(segments_path.as_path())
.expect(&format!("Can't read segments directory {}", segments_path.to_str().unwrap()));
@@ -44,28 +39,26 @@ pub fn load_collection(
let segments_path = entry.unwrap().path();
let segment = match load_segment(segments_path.as_path()) {
Ok(x) => x,
Err(err) => panic!(
format!("Can't load segments from {}, error: {}", segments_path.to_str().unwrap(), err)
),
Err(err) => panic!("Can't load segments from {}, error: {}", segments_path.to_str().unwrap(), err),
};
segment_holder.add(segment);
};
let segment_config = load_config(&collection_path);
let optimizers = build_optimizers(
collection_path,
&segment_config,
&optimizers_config,
&collection_config.params,
&collection_config.optimizer_config,
&collection_config.hnsw_config,
);
let collection = construct_collection(
segment_holder,
&segment_config,
collection_config,
wal,
search_runtime,
optimizers,
optimizers_config.flush_interval_sec,
collection_path
);
{
@@ -78,7 +71,7 @@ pub fn load_collection(
match collection.updater.update(op_num, update) {
Ok(_) => {}
Err(err) => match err {
CollectionError::ServiceError { error } => panic!(format!("Can't apply WAL operation: {}", error)),
CollectionError::ServiceError { error } => panic!("Can't apply WAL operation: {}", error),
_ => {}
}
}

24
lib/collection/src/collection_builder/optimizers_builder.rs
View File

@@ -1,31 +1,42 @@
use crate::update_handler::update_handler::Optimizer;
use std::sync::Arc;
use crate::segment_manager::optimizers::vacuum_optimizer::VacuumOptimizer;
use segment::types::SegmentConfig;
use segment::types::{HnswConfig};
use crate::segment_manager::optimizers::merge_optimizer::MergeOptimizer;
use std::path::Path;
use serde::{Deserialize, Serialize};
use schemars::{JsonSchema};
use crate::segment_manager::optimizers::indexing_optimizer::IndexingOptimizer;
use crate::segment_manager::optimizers::segment_optimizer::OptimizerThresholds;
use crate::config::CollectionParams;
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
pub struct OptimizersConfig {
/// The minimal fraction of deleted vectors in a segment, required to perform segment optimization
pub deleted_threshold: f64,
/// The minimal number of vectors in a segment, required to perform segment optimization
pub vacuum_min_vector_number: usize,
/// If the number of segments exceeds this value, the optimizer will merge the smallest segments.
pub max_segment_number: usize,
/// Maximum number of vectors to store in-memory per segment.
/// Segments larger than this threshold will be stored as read-only memmaped file.
pub memmap_threshold: usize,
/// Maximum number of vectors allowed for plain index.
/// Default value based on https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md
pub indexing_threshold: usize,
/// Starting from this amount of vectors per-segment the engine will start building index for payload.
pub payload_indexing_threshold: usize,
/// Minimum interval between forced flushes.
pub flush_interval_sec: u64,
}
pub fn build_optimizers(
collection_path: &Path,
segment_config: &SegmentConfig,
collection_params: &CollectionParams,
optimizers_config: &OptimizersConfig,
hnsw_config: &HnswConfig,
) -> Arc<Vec<Box<Optimizer>>> {
let segments_path = collection_path.join("segments");
let temp_segments_path = collection_path.join("temp_segments");
@@ -42,7 +53,8 @@ pub fn build_optimizers(
threshold_config.clone(),
segments_path.clone(),
temp_segments_path.clone(),
segment_config.clone(),
collection_params.clone(),
hnsw_config.clone(),
)
),
Box::new(
@@ -51,7 +63,8 @@ pub fn build_optimizers(
threshold_config.clone(),
segments_path.clone(),
temp_segments_path.clone(),
segment_config.clone(),
collection_params.clone(),
hnsw_config.clone(),
)
),
Box::new(VacuumOptimizer::new(
@@ -60,7 +73,8 @@ pub fn build_optimizers(
threshold_config.clone(),
segments_path.clone(),
temp_segments_path.clone(),
segment_config.clone(),
collection_params.clone(),
hnsw_config.clone(),
))
])
}

80
lib/collection/src/config.rs Normal file
View File

@@ -0,0 +1,80 @@
use std::fs::File;
use std::io::{Read, Write};
use std::path::Path;
use atomicwrites::AtomicFile;
use atomicwrites::OverwriteBehavior::AllowOverwrite;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use wal::WalOptions;
use segment::types::{Distance, HnswConfig};
use crate::collection_builder::optimizers_builder::OptimizersConfig;
use crate::operations::types::{CollectionError, CollectionResult};
pub const COLLECTION_CONFIG_FILE: &str = "config.json";
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
pub struct WalConfig {
/// Size of a single WAL segment in MB
pub wal_capacity_mb: usize,
/// Number of WAL segments to create ahead of actually used ones
pub wal_segments_ahead: usize,
}
impl From<&WalConfig> for WalOptions {
fn from(config: &WalConfig) -> Self {
WalOptions {
segment_capacity: config.wal_capacity_mb * 1024 * 1024,
segment_queue_len: config.wal_segments_ahead
}
}
}
impl Default for WalConfig {
fn default() -> Self { WalConfig { wal_capacity_mb: 32, wal_segments_ahead: 0 } }
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
#[serde(rename_all = "snake_case")]
pub struct CollectionParams {
/// Size of a vectors used
pub vector_size: usize,
/// Type of distance function used for measuring distance between vectors
pub distance: Distance
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
pub struct CollectionConfig {
pub params: CollectionParams,
pub hnsw_config: HnswConfig,
pub optimizer_config: OptimizersConfig,
pub wal_config: WalConfig,
}
impl CollectionConfig {
pub fn save(&self, path: &Path) -> CollectionResult<()> {
let config_path = path.join(COLLECTION_CONFIG_FILE);
let af = AtomicFile::new(&config_path, AllowOverwrite);
let state_bytes = serde_json::to_vec(self).unwrap();
af.write(|f| {
f.write_all(&state_bytes)
}).or_else(move |err|
Err(CollectionError::ServiceError {
error: format!("Can't write {:?}, error: {}", config_path, err)
})
)?;
Ok(())
}
pub fn load(path: &Path) -> CollectionResult<Self> {
let config_path = path.join(COLLECTION_CONFIG_FILE);
let mut contents = String::new();
let mut file = File::open(config_path)?;
file.read_to_string(&mut contents)?;
Ok(serde_json::from_str(&contents)?)
}
}

1
lib/collection/src/lib.rs
View File

@@ -1,4 +1,5 @@
pub mod collection_builder;
pub mod config;
mod update_handler;
pub mod operations;
pub mod collection;

118
lib/collection/src/operations/config_diff.rs Normal file
View File

@@ -0,0 +1,118 @@
use serde::{Deserialize, Serialize};
use merge::Merge;
use serde::de::DeserializeOwned;
use schemars::{JsonSchema};
use crate::operations::types::CollectionResult;
use segment::types::HnswConfig;
use crate::config::WalConfig;
use crate::collection_builder::optimizers_builder::OptimizersConfig;
// Structures for partial update of collection params
// ToDo: Make auto-generated somehow...
pub trait DiffConfig<T: DeserializeOwned + Serialize> {
fn update(self, config: &T) -> CollectionResult<T>
where Self: Sized,
Self: Serialize,
Self: DeserializeOwned,
Self: Merge { update_config(config, self) }
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Copy, Clone, PartialEq, Eq, Merge)]
#[serde(rename_all = "snake_case")]
pub struct HnswConfigDiff {
/// Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.
pub m: Option<usize>,
/// Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.
pub ef_construct: Option<usize>,
/// Minimal amount of points for additional payload-based indexing.
/// If payload chunk is smaller than `full_scan_threshold` additional indexing won't be used -
/// in this case full-scan search should be preferred by query planner and additional indexing is not required.
pub full_scan_threshold: Option<usize>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Merge)]
pub struct WalConfigDiff {
/// Size of a single WAL segment in MB
pub wal_capacity_mb: Option<usize>,
/// Number of WAL segments to create ahead of actually used ones
pub wal_segments_ahead: Option<usize>,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Merge)]
pub struct OptimizersConfigDiff {
/// The minimal fraction of deleted vectors in a segment, required to perform segment optimization
pub deleted_threshold: Option<f64>,
/// The minimal number of vectors in a segment, required to perform segment optimization
pub vacuum_min_vector_number: Option<usize>,
/// If the number of segments exceeds this value, the optimizer will merge the smallest segments.
pub max_segment_number: Option<usize>,
/// Maximum number of vectors to store in-memory per segment.
/// Segments larger than this threshold will be stored as read-only memmaped file.
pub memmap_threshold: Option<usize>,
/// Maximum number of vectors allowed for plain index.
/// Default value based on https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md
pub indexing_threshold: Option<usize>,
/// Starting from this amount of vectors per-segment the engine will start building index for payload.
pub payload_indexing_threshold: Option<usize>,
/// Minimum interval between forced flushes.
pub flush_interval_sec: Option<u64>,
}
impl DiffConfig<HnswConfig> for HnswConfigDiff {}
impl DiffConfig<OptimizersConfig> for OptimizersConfigDiff {}
impl DiffConfig<WalConfig> for WalConfigDiff {}
/// Hacky way to update configuration structures with diff-updates.
/// Intended to only be used in non critical for speed places.
/// ToDo: Replace with proc macro
pub fn update_config<T: DeserializeOwned + Serialize, Y: DeserializeOwned + Serialize + Merge>(config: &T, mut update: Y) -> CollectionResult<T> {
let serialized = serde_json::to_vec(config)?;
let config_as_diff: Y = serde_json::from_slice(serialized.as_slice())?;
update.merge(config_as_diff);
let serialized = serde_json::to_vec(&update)?;
let res = serde_json::from_slice(serialized.as_slice())?;
Ok(res)
}
#[cfg(test)]
mod tests {
use super::*;
use segment::types::HnswConfig;
use crate::collection_builder::optimizers_builder::OptimizersConfig;
#[test]
fn test_hnsw_update() {
let base_config = HnswConfig::default();
let update: HnswConfigDiff = serde_json::from_str(&r#"{ "m": 32 }"#).unwrap();
let new_config = update.update(&base_config).unwrap();
assert_eq!(new_config.m, 32)
}
#[test]
fn test_optimizer_update() {
let base_config = OptimizersConfig {
deleted_threshold: 0.9,
vacuum_min_vector_number: 1000,
max_segment_number: 10,
memmap_threshold: 100_000,
indexing_threshold: 50_000,
payload_indexing_threshold: 20_000,
flush_interval_sec: 30,
};
let update: OptimizersConfigDiff = serde_json::from_str(&r#"{ "indexing_threshold": 10000 }"#).unwrap();
let new_config = update.update(&base_config).unwrap();
assert_eq!(new_config.indexing_threshold, 10000)
}
#[test]
fn test_wal_config() {
let base_config = WalConfig::default();
let update: WalConfigDiff = serde_json::from_str(&r#"{ "wal_segments_ahead": 2 }"#).unwrap();
let new_config = update.update(&base_config).unwrap();
assert_eq!(new_config.wal_segments_ahead, 2)
}
}

1
lib/collection/src/operations/mod.rs
View File

@@ -1,6 +1,7 @@
pub mod types;
pub mod point_ops;
pub mod payload_ops;
pub mod config_diff;
use serde::{Deserialize, Serialize};
use schemars::{JsonSchema};

86
lib/collection/src/operations/types.rs
View File

@@ -1,11 +1,32 @@
use segment::types::{VectorElementType, PointIdType, TheMap, PayloadKeyType, PayloadType, SeqNumberType, Filter, SearchParams, SegmentConfig};
use crossbeam_channel::SendError;
use futures::io;
use schemars::JsonSchema;
use serde;
use serde::{Deserialize, Serialize};
use schemars::{JsonSchema};
use thiserror::Error;
use serde_json::Error as JsonError;
use tokio::task::JoinError;
use std::result;
use segment::entry::entry_point::OperationError;
use segment::types::{Filter, PayloadKeyType, PayloadType, PointIdType, SearchParams, SeqNumberType, TheMap, VectorElementType};
use crate::config::CollectionConfig;
use crate::wal::WalError;
/// Type of vector in API
pub type VectorType = Vec<VectorElementType>;
#[derive(Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum CollectionStatus {
/// Collection if completely ready for requests
Green,
/// Collection is available, but some segments might be under optimization
Yellow,
/// Something is not OK
Red
}
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
@@ -22,6 +43,8 @@ pub struct Record {
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
/// Current statistics and configuration of the collection.
pub struct CollectionInfo {
/// Status of the collection
pub status: CollectionStatus,
/// Number of vectors in collection
pub vectors_count: usize,
/// Number of segments in collection
@@ -31,7 +54,7 @@ pub struct CollectionInfo {
/// RAM used by collection
pub ram_data_size: usize,
/// Collection settings
pub config: SegmentConfig,
pub config: CollectionConfig,
}
@@ -86,3 +109,60 @@ pub struct RecommendRequest {
}
#[derive(Error, Debug, Clone)]
#[error("{0}")]
pub enum CollectionError {
#[error("Wrong input: {description}")]
BadInput { description: String },
#[error("No point with id {missed_point_id} found")]
NotFound { missed_point_id: PointIdType },
#[error("Service internal error: {error}")]
ServiceError { error: String },
#[error("Bad request: {description}")]
BadRequest { description: String },
}
impl From<OperationError> for CollectionError {
fn from(err: OperationError) -> Self {
match err {
OperationError::WrongVector { .. } => Self::BadInput { description: format!("{}", err) },
OperationError::PointIdError { missed_point_id } => Self::NotFound { missed_point_id },
OperationError::ServiceError { description } => Self::ServiceError { error: description },
OperationError::TypeError { .. } => Self::BadInput { description: format!("{}", err) },
}
}
}
impl From<JoinError> for CollectionError {
fn from(err: JoinError) -> Self {
Self::ServiceError { error: format!("{}", err) }
}
}
impl From<WalError> for CollectionError {
fn from(err: WalError) -> Self {
Self::ServiceError { error: format!("{}", err) }
}
}
impl<T> From<SendError<T>> for CollectionError {
fn from(_err: SendError<T>) -> Self {
Self::ServiceError { error: format!("Can't reach one of the workers") }
}
}
impl From<JsonError> for CollectionError {
fn from(err: JsonError) -> Self {
CollectionError::ServiceError { error: format!("Json error: {}", err) }
}
}
impl From<io::Error> for CollectionError {
fn from(err: io::Error) -> Self {
CollectionError::ServiceError { error: format!("File IO error: {}", err) }
}
}
pub type CollectionResult<T> = result::Result<T, CollectionError>;

2
lib/collection/src/segment_manager/holders/proxy_segment.rs
View File

@@ -328,7 +328,7 @@ mod tests {
let mut seen_points: HashSet<PointIdType> = Default::default();
for res in search_result {
if seen_points.contains(&res.id) {
assert!(false, format!("point {} appears multiple times", res.id));
assert!(false, "point {} appears multiple times", res.id);
}
seen_points.insert(res.id);
}

29
lib/collection/src/segment_manager/optimizers/indexing_optimizer.rs
View File

@@ -1,15 +1,17 @@
use std::path::{PathBuf, Path};
use segment::types::{SegmentConfig, SegmentType};
use segment::types::{SegmentType, HnswConfig};
use crate::segment_manager::holders::segment_holder::{LockedSegmentHolder, SegmentId, LockedSegment};
use std::cmp::min;
use crate::segment_manager::optimizers::segment_optimizer::{SegmentOptimizer, OptimizerThresholds};
use crate::config::CollectionParams;
pub struct IndexingOptimizer {
thresholds_config: OptimizerThresholds,
segments_path: PathBuf,
collection_temp_dir: PathBuf,
config: SegmentConfig,
collection_params: CollectionParams,
hnsw_config: HnswConfig,
}
impl IndexingOptimizer {
@@ -17,13 +19,15 @@ impl IndexingOptimizer {
thresholds_config: OptimizerThresholds,
segments_path: PathBuf,
collection_temp_dir: PathBuf,
config: SegmentConfig,
collection_params: CollectionParams,
hnsw_config: HnswConfig,
) -> Self {
IndexingOptimizer {
thresholds_config,
segments_path,
collection_temp_dir,
config,
collection_params,
hnsw_config,
}
}
@@ -61,8 +65,12 @@ impl SegmentOptimizer for IndexingOptimizer {
self.collection_temp_dir.as_path()
}
fn base_segment_config(&self) -> SegmentConfig {
self.config.clone()
fn collection_params(&self) -> CollectionParams {
self.collection_params.clone()
}
fn hnsw_config(&self) -> HnswConfig {
self.hnsw_config.clone()
}
fn threshold_config(&self) -> &OptimizerThresholds {
@@ -128,13 +136,11 @@ mod tests {
},
segments_dir.path().to_owned(),
segments_temp_dir.path().to_owned(),
SegmentConfig {
CollectionParams {
vector_size: segment_config.vector_size,
index: Default::default(),
payload_index: Some(Default::default()),
distance: segment_config.distance,
storage_type: StorageType::default(),
},
Default::default()
);
let locked_holder = Arc::new(RwLock::new(holder));
@@ -156,7 +162,10 @@ mod tests {
// ------ Plain -> Mmap & Indexed payload
let suggested_to_optimize = index_optimizer.check_condition(locked_holder.clone());
assert!(suggested_to_optimize.contains(&large_segment_id));
eprintln!("suggested_to_optimize = {:#?}", suggested_to_optimize);
index_optimizer.optimize(locked_holder.clone(), suggested_to_optimize).unwrap();
eprintln!("Done");
// ------ Plain -> Indexed payload
let suggested_to_optimize = index_optimizer.check_condition(locked_holder.clone());

33
lib/collection/src/segment_manager/optimizers/merge_optimizer.rs
View File

@@ -1,8 +1,9 @@
use crate::segment_manager::optimizers::segment_optimizer::{SegmentOptimizer, OptimizerThresholds};
use crate::segment_manager::holders::segment_holder::{LockedSegmentHolder, SegmentId};
use segment::types::{SegmentType, SegmentConfig};
use segment::types::{SegmentType, HnswConfig};
use itertools::Itertools;
use std::path::{PathBuf, Path};
use crate::config::CollectionParams;
/// Optimizer that tries to reduce number of segments until it fits configured value
@@ -11,7 +12,8 @@ pub struct MergeOptimizer {
thresholds_config: OptimizerThresholds,
segments_path: PathBuf,
collection_temp_dir: PathBuf,
config: SegmentConfig,
collection_params: CollectionParams,
hnsw_config: HnswConfig,
}
impl MergeOptimizer {
@@ -20,13 +22,16 @@ impl MergeOptimizer {
thresholds_config: OptimizerThresholds,
segments_path: PathBuf,
collection_temp_dir: PathBuf,
config: SegmentConfig) -> Self {
collection_params: CollectionParams,
hnsw_config: HnswConfig,
) -> Self {
return MergeOptimizer {
max_segments,
thresholds_config,
segments_path,
collection_temp_dir,
config,
collection_params,
hnsw_config,
};
}
}
@@ -41,15 +46,18 @@ impl SegmentOptimizer for MergeOptimizer {
self.collection_temp_dir.as_path()
}
fn base_segment_config(&self) -> SegmentConfig {
self.config.clone()
fn collection_params(&self) -> CollectionParams {
self.collection_params.clone()
}
fn hnsw_config(&self) -> HnswConfig {
self.hnsw_config.clone()
}
fn threshold_config(&self) -> &OptimizerThresholds {
&self.thresholds_config
}
fn check_condition(&self, segments: LockedSegmentHolder) -> Vec<SegmentId> {
let read_segments = segments.read();
@@ -82,7 +90,7 @@ mod tests {
use super::*;
use crate::segment_manager::fixtures::{random_segment};
use crate::segment_manager::holders::segment_holder::{SegmentHolder, LockedSegment};
use segment::types::{Distance, Indexes};
use segment::types::{Distance};
use std::sync::{Arc};
use tempdir::TempDir;
use parking_lot::RwLock;
@@ -119,13 +127,12 @@ mod tests {
},
dir.path().to_owned(),
temp_dir.path().to_owned(),
SegmentConfig {
CollectionParams {
vector_size: 4,
index: Indexes::Plain {},
payload_index: Some(Default::default()),
distance: Distance::Dot,
storage_type: Default::default(),
});
},
Default::default()
);
let locked_holder = Arc::new(RwLock::new(holder));

53
lib/collection/src/segment_manager/optimizers/segment_optimizer.rs
View File

@@ -1,5 +1,5 @@
use segment::types::{PointIdType, PayloadKeyType, SegmentConfig, Indexes, StorageType, PayloadIndexType};
use crate::collection::CollectionResult;
use segment::types::{PointIdType, PayloadKeyType, SegmentConfig, Indexes, StorageType, PayloadIndexType, HnswConfig};
use crate::operations::types::CollectionResult;
use crate::segment_manager::holders::segment_holder::{SegmentId, LockedSegment, LockedSegmentHolder};
use std::sync::Arc;
use segment::segment::Segment;
@@ -12,6 +12,7 @@ use segment::segment_constructor::segment_builder::SegmentBuilder;
use std::convert::TryInto;
use std::path::Path;
use segment::segment_constructor::simple_segment_constructor::build_simple_segment;
use crate::config::CollectionParams;
#[derive(Debug, Clone)]
@@ -29,7 +30,10 @@ pub trait SegmentOptimizer {
fn temp_path(&self) -> &Path;
/// Get basic segment config
fn base_segment_config(&self) -> SegmentConfig;
fn collection_params(&self) -> CollectionParams;
/// Get HNSW config
fn hnsw_config(&self) -> HnswConfig;
/// Get thresholds configuration for the current optimizer
fn threshold_config(&self) -> &OptimizerThresholds;
@@ -39,7 +43,14 @@ pub trait SegmentOptimizer {
/// Build temp segment
fn temp_segment(&self) -> CollectionResult<LockedSegment> {
let config = self.base_segment_config();
let collection_params = self.collection_params();
let config = SegmentConfig {
vector_size: collection_params.vector_size,
distance: collection_params.distance,
index: Indexes::Plain {},
payload_index: Some(PayloadIndexType::Plain),
storage_type: StorageType::InMemory
};
Ok(LockedSegment::new(build_simple_segment(
self.collection_path(),
config.vector_size,
@@ -55,31 +66,23 @@ pub trait SegmentOptimizer {
let have_indexed_fields = optimizing_segments.iter()
.any(|s| !s.get().read().get_indexed_fields().is_empty());
let mut optimized_config = self.base_segment_config();
let thresholds = self.threshold_config();
let collection_params = self.collection_params();
if total_vectors < thresholds.memmap_threshold {
optimized_config.storage_type = StorageType::InMemory;
} else {
optimized_config.storage_type = StorageType::Mmap;
}
if total_vectors < thresholds.indexing_threshold {
optimized_config.index = Indexes::Plain {};
} else {
optimized_config.index = match optimized_config.index {
Indexes::Plain { } => Indexes::default_hnsw(),
_ => optimized_config.index
}
}
let is_indexed = total_vectors >= thresholds.indexing_threshold;
// Create structure index only if there is something to index
if total_vectors < thresholds.payload_indexing_threshold || !have_indexed_fields {
optimized_config.payload_index = Some(PayloadIndexType::Plain)
} else {
optimized_config.payload_index = Some(PayloadIndexType::Struct);
}
let is_payload_indexed = total_vectors >= thresholds.payload_indexing_threshold && have_indexed_fields;
let is_on_disk = total_vectors >= thresholds.memmap_threshold;
let optimized_config = SegmentConfig {
vector_size: collection_params.vector_size,
distance: collection_params.distance,
index: if is_indexed { Indexes::Hnsw(self.hnsw_config()) } else { Indexes::Plain {} },
payload_index: Some(if is_payload_indexed { PayloadIndexType::Struct } else { PayloadIndexType::Plain }),
storage_type: if is_on_disk { StorageType::Mmap } else { StorageType::InMemory }
};
Ok(SegmentBuilder::new(
self.collection_path(),

33
lib/collection/src/segment_manager/optimizers/vacuum_optimizer.rs
View File

@@ -1,8 +1,9 @@
use crate::segment_manager::holders::segment_holder::{SegmentId, LockedSegment, LockedSegmentHolder};
use segment::types::{SegmentType, SegmentConfig};
use segment::types::{SegmentType, HnswConfig};
use ordered_float::OrderedFloat;
use crate::segment_manager::optimizers::segment_optimizer::{SegmentOptimizer, OptimizerThresholds};
use std::path::{PathBuf, Path};
use crate::config::CollectionParams;
pub struct VacuumOptimizer {
@@ -11,7 +12,8 @@ pub struct VacuumOptimizer {
thresholds_config: OptimizerThresholds,
segments_path: PathBuf,
collection_temp_dir: PathBuf,
config: SegmentConfig,
collection_params: CollectionParams,
hnsw_config: HnswConfig,
}
@@ -21,14 +23,17 @@ impl VacuumOptimizer {
thresholds_config: OptimizerThresholds,
segments_path: PathBuf,
collection_temp_dir: PathBuf,
config: SegmentConfig) -> Self {
collection_params: CollectionParams,
hnsw_config: HnswConfig,
) -> Self {
VacuumOptimizer {
deleted_threshold,
min_vectors_number,
thresholds_config,
segments_path,
collection_temp_dir,
config,
collection_params,
hnsw_config
}
}
@@ -64,8 +69,12 @@ impl SegmentOptimizer for VacuumOptimizer {
self.collection_temp_dir.as_path()
}
fn base_segment_config(&self) -> SegmentConfig {
self.config.clone()
fn collection_params(&self) -> CollectionParams {
self.collection_params.clone()
}
fn hnsw_config(&self) -> HnswConfig {
self.hnsw_config.clone()
}
fn threshold_config(&self) -> &OptimizerThresholds {
@@ -89,7 +98,7 @@ mod tests {
use itertools::Itertools;
use rand::Rng;
use std::sync::Arc;
use segment::types::{Distance, Indexes, PayloadType, StorageType};
use segment::types::{Distance, PayloadType};
use tempdir::TempDir;
use parking_lot::RwLock;
@@ -154,20 +163,18 @@ mod tests {
let vacuum_optimizer = VacuumOptimizer::new(
0.2,
50,
OptimizerThresholds{
OptimizerThresholds {
memmap_threshold: 1000000,
indexing_threshold: 1000000,
payload_indexing_threshold: 1000000
payload_indexing_threshold: 1000000,
},
dir.path().to_owned(),
temp_dir.path().to_owned(),
SegmentConfig {
CollectionParams {
vector_size: 4,
index: Indexes::Plain {},
payload_index: Some(Default::default()),
distance: Distance::Dot,
storage_type: StorageType::InMemory,
},
Default::default()
);
let suggested_to_optimize = vacuum_optimizer.check_condition(locked_holder.clone());

9
lib/collection/src/segment_manager/segment_managers.rs
View File

@@ -1,9 +1,10 @@
use segment::types::{SeqNumberType, ScoredPoint, PointIdType};
use crate::collection::{CollectionResult};
use crate::operations::CollectionUpdateOperations;
use crate::operations::types::{Record, SearchRequest};
use std::sync::Arc;
use segment::types::{PointIdType, ScoredPoint, SeqNumberType};
use crate::operations::CollectionUpdateOperations;
use crate::operations::types::{CollectionResult, Record, SearchRequest};
pub trait SegmentSearcher {
fn search(&self,
// Request is supposed to be a read only, that is why no mutex used

5
lib/collection/src/segment_manager/simple_segment_searcher.rs
View File

@@ -1,7 +1,7 @@
use crate::segment_manager::holders::segment_holder::{LockedSegment, LockedSegmentHolder};
use std::sync::Arc;
use crate::segment_manager::segment_managers::{SegmentSearcher};
use crate::collection::CollectionResult;
use crate::operations::types::CollectionResult;
use segment::types::{ScoredPoint, PointIdType, SeqNumberType};
use tokio::runtime::Runtime;
use std::collections::{HashSet, HashMap};
@@ -55,8 +55,6 @@ impl SegmentSearcher for SimpleSegmentSearcher {
return Ok(vec![]);
}
let distance = some_segment.unwrap().1.get().read().config().distance;
let searches: Vec<_> = segments
.iter()
.map(|(_id, segment)|
@@ -89,7 +87,6 @@ impl SegmentSearcher for SimpleSegmentSearcher {
!res
}),
request.top,
&distance,
);
Ok(top_scores)

33
lib/collection/src/segment_manager/simple_segment_updater.rs
View File

@@ -1,25 +1,24 @@
use std::sync::Mutex;
use crate::segment_manager::holders::segment_holder::{LockedSegmentHolder};
use crate::segment_manager::segment_managers::SegmentUpdater;
use crate::operations::{CollectionUpdateOperations, FieldIndexOperations};
use crate::collection::{CollectionResult, CollectionError};
use segment::types::{SeqNumberType, PointIdType, PayloadKeyType, PayloadInterface};
use std::collections::{HashSet, HashMap};
use crate::operations::types::VectorType;
use std::collections::{HashMap, HashSet};
use crate::operations::point_ops::{PointOperations, PointInsertOperations};
use segment::types::{PayloadKeyType, PointIdType, SeqNumberType, PayloadInterface};
use crate::operations::{CollectionUpdateOperations, FieldIndexOperations};
use crate::operations::point_ops::{PointInsertOperations, PointOperations};
use crate::operations::types::{CollectionError, CollectionResult, VectorType};
use crate::segment_manager::holders::segment_holder::LockedSegmentHolder;
use crate::segment_manager::segment_managers::SegmentUpdater;
use crate::operations::payload_ops::PayloadOps;
pub struct SimpleSegmentUpdater {
segments: LockedSegmentHolder,
update_lock: Mutex<bool>,
// update_lock: Mutex<bool>,
}
impl SimpleSegmentUpdater {
pub fn new(segments: LockedSegmentHolder) -> Self {
SimpleSegmentUpdater {
segments,
update_lock: Mutex::new(false),
// update_lock: Mutex::new(false),
}
}
@@ -279,12 +278,14 @@ impl SegmentUpdater for SimpleSegmentUpdater {
#[cfg(test)]
mod tests {
use super::*;
use crate::segment_manager::fixtures::{build_searcher};
use crate::segment_manager::segment_managers::SegmentSearcher;
use segment::types::PayloadVariant;
use tempdir::TempDir;
use segment::types::PayloadVariant;
use crate::segment_manager::fixtures::build_searcher;
use crate::segment_manager::segment_managers::SegmentSearcher;
use super::*;
#[test]
fn test_point_ops() {
let dir = TempDir::new("segment_dir").unwrap();
@@ -293,7 +294,6 @@ mod tests {
let updater = SimpleSegmentUpdater {
segments: searcher.segments.clone(),
update_lock: Mutex::new(false),
};
let points = vec![1, 500];
@@ -349,7 +349,6 @@ mod tests {
let updater = SimpleSegmentUpdater {
segments: searcher.segments.clone(),
update_lock: Mutex::new(false),
};
let mut payload: HashMap<PayloadKeyType, PayloadInterface> = Default::default();

50
lib/collection/src/update_handler/update_handler.rs
View File

@@ -10,22 +10,27 @@ use crate::operations::CollectionUpdateOperations;
use tokio::time::{Duration, Instant};
use tokio::runtime::Runtime;
use log::debug;
use crate::operations::types::CollectionResult;
pub type Optimizer = dyn SegmentOptimizer + Sync + Send;
pub enum UpdateSignal {
/// Info that operation with
Operation(SeqNumberType),
/// Stop all optimizers and listening
Stop,
/// Empty signal used to trigger optimizers
Nop,
}
pub struct UpdateHandler {
optimizers: Arc<Vec<Box<Optimizer>>>,
pub optimizers: Arc<Vec<Box<Optimizer>>>,
pub flush_timeout_sec: u64,
segments: LockedSegmentHolder,
receiver: Receiver<UpdateSignal>,
worker: Option<JoinHandle<()>>,
runtime_handle: Arc<Runtime>,
wal: Arc<Mutex<SerdeWal<CollectionUpdateOperations>>>,
flush_timeout_sec: u64,
}
@@ -63,6 +68,35 @@ impl UpdateHandler {
));
}
/// Gracefully wait before all optimizations stop.
/// If some optimization is in progress - it will be finished before shutdown.
/// Blocking function.
pub fn wait_worker_stops(&mut self) -> CollectionResult<()> {
let res = match &mut self.worker {
None => (),
Some(handle) => self.runtime_handle.block_on(handle)?
};
self.worker = None;
Ok(res)
}
fn process_optimization(
optimizers: Arc<Vec<Box<Optimizer>>>,
segments: LockedSegmentHolder,
) {
for optimizer in optimizers.iter() {
let mut nonoptimal_segment_ids = optimizer.check_condition(segments.clone());
while !nonoptimal_segment_ids.is_empty() {
debug!("Start optimization on segments: {:?}", nonoptimal_segment_ids);
// If optimization fails, it could not be reported to anywhere except for console.
// So the only recovery here is to stop optimization and await for restart
optimizer.optimize(segments.clone(), nonoptimal_segment_ids).unwrap();
nonoptimal_segment_ids = optimizer.check_condition(segments.clone());
}
}
}
async fn worker_fn(
optimizers: Arc<Vec<Box<Optimizer>>>,
receiver: Receiver<UpdateSignal>,
@@ -77,15 +111,13 @@ impl UpdateHandler {
match recv_res {
Ok(signal) => {
match signal {
UpdateSignal::Nop => {
Self::process_optimization(optimizers.clone(), segments.clone());
}
UpdateSignal::Operation(operation_id) => {
debug!("Performing update operation: {}", operation_id);
for optimizer in optimizers.iter() {
let unoptimal_segment_ids = optimizer.check_condition(segments.clone());
if !unoptimal_segment_ids.is_empty() {
debug!("Start optimization on segments: {:?}", unoptimal_segment_ids);
optimizer.optimize(segments.clone(), unoptimal_segment_ids).unwrap();
}
}
Self::process_optimization(optimizers.clone(), segments.clone());
let elapsed = last_flushed.elapsed();
if elapsed > flush_timeout {
debug!("Performing flushing: {}", operation_id);

19
lib/collection/tests/collection_test.rs
View File

@@ -3,14 +3,13 @@ mod common;
use collection::operations::CollectionUpdateOperations;
use collection::operations::point_ops::{PointOperations, PointStruct};
use crate::common::{simple_collection_fixture, TEST_OPTIMIZERS_CONFIG};
use crate::common::{simple_collection_fixture};
use collection::operations::types::{UpdateStatus, SearchRequest, RecommendRequest};
use std::sync::Arc;
use collection::operations::payload_ops::PayloadOps;
use std::collections::HashMap;
use segment::types::{PayloadKeyType, PayloadVariant, PayloadInterface};
use collection::collection_builder::collection_loader::load_collection;
use wal::WalOptions;
use tempdir::TempDir;
use tokio::runtime;
use collection::operations::point_ops::PointInsertOperations::{BatchPoints, PointsList};
@@ -42,7 +41,7 @@ fn test_collection_updater() {
Ok(res) => {
assert_eq!(res.status, UpdateStatus::Completed)
}
Err(err) => assert!(false, format!("operation failed: {:?}", err)),
Err(err) => assert!(false, "operation failed: {:?}", err),
}
let search_request = Arc::new(SearchRequest {
@@ -60,7 +59,7 @@ fn test_collection_updater() {
assert_eq!(res.len(), 3);
assert_eq!(res[0].id, 2);
}
Err(err) => assert!(false, format!("search failed: {:?}", err)),
Err(err) => assert!(false, "search failed: {:?}", err),
}
}
@@ -105,12 +104,6 @@ fn test_collection_loading() {
collection.update(assign_payload, true).unwrap();
}
let wal_options = WalOptions {
segment_capacity: 100,
segment_queue_len: 0,
};
let rt = Arc::new(runtime::Builder::new_multi_thread()
.max_threads(2)
.build().unwrap());
@@ -120,9 +113,7 @@ fn test_collection_loading() {
let loaded_collection = load_collection(
collection_dir.path(),
&wal_options,
rt.clone(),
&TEST_OPTIMIZERS_CONFIG,
);
let retrieved = loaded_collection.retrieve(&vec![1, 2], true, true).unwrap();
@@ -195,9 +186,9 @@ fn test_deserialization2() {
eprintln!("read_obj = {:#?}", read_obj);
let crob_bytes = rmp_serde::to_vec(&insert_points).unwrap();
let raw_bytes = rmp_serde::to_vec(&insert_points).unwrap();
let read_obj2: CollectionUpdateOperations = rmp_serde::from_read_ref(&crob_bytes).unwrap();
let read_obj2: CollectionUpdateOperations = rmp_serde::from_read_ref(&raw_bytes).unwrap();
eprintln!("read_obj2 = {:#?}", read_obj2);
}

30
lib/collection/tests/common/mod.rs
View File

@@ -1,13 +1,13 @@
use collection::collection_builder::collection_builder::build_collection;
use wal::WalOptions;
use collection::collection::Collection;
use segment::types::{Distance, SegmentConfig, Indexes};
use segment::types::{Distance};
use tokio::runtime::Runtime;
use tokio::runtime;
use std::sync::Arc;
use std::path::Path;
use collection::collection_builder::optimizers_builder::OptimizersConfig;
use collection::collection_builder::collection_loader::load_collection;
use collection::config::{WalConfig, CollectionParams};
pub const TEST_OPTIMIZERS_CONFIG: OptimizersConfig = OptimizersConfig {
@@ -23,11 +23,6 @@ pub const TEST_OPTIMIZERS_CONFIG: OptimizersConfig = OptimizersConfig {
#[allow(dead_code)]
pub fn load_collection_fixture(collection_path: &Path) -> (Arc<Runtime>, Collection) {
let wal_options = WalOptions {
segment_capacity: 100,
segment_queue_len: 0,
};
let threaded_rt = Arc::new(runtime::Builder::new_multi_thread()
.max_threads(2)
.build().unwrap());
@@ -35,29 +30,21 @@ pub fn load_collection_fixture(collection_path: &Path) -> (Arc<Runtime>, Collect
let collection = load_collection(
collection_path,
&wal_options,
threaded_rt.clone(),
&TEST_OPTIMIZERS_CONFIG,
);
return (threaded_rt, collection);
}
pub fn simple_collection_fixture(collection_path: &Path) -> (Arc<Runtime>, Collection) {
let wal_options = WalOptions {
segment_capacity: 100,
segment_queue_len: 0,
let wal_config = WalConfig {
wal_capacity_mb: 1,
wal_segments_ahead: 0
};
let collection_config = SegmentConfig {
let collection_params = CollectionParams {
vector_size: 4,
index: Indexes::Hnsw {
m: 16,
ef_construct: 128,
},
payload_index: Some(Default::default()),
distance: Distance::Dot,
storage_type: Default::default(),
};
let threaded_rt = Arc::new(runtime::Builder::new_multi_thread()
@@ -67,10 +54,11 @@ pub fn simple_collection_fixture(collection_path: &Path) -> (Arc<Runtime>, Colle
let collection = build_collection(
collection_path,
&wal_options,
&collection_config,
&wal_config,
&collection_params,
threaded_rt.clone(),
&TEST_OPTIMIZERS_CONFIG,
&Default::default()
).unwrap();
return (threaded_rt, collection);

23
lib/segment/Cargo.toml
View File

@@ -1,23 +1,23 @@
[package]
name = "segment"
version = "0.2.1"
version = "0.3.0"
authors = ["Andrey Vasnetsov <vasnetsov93@gmail.com>"]
edition = "2018"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dev-dependencies]
pprof = { version = "0.4", features = ["flamegraph"] }
tempdir = "0.3.7"
criterion = "0.3"
rand = "0.8"
[dependencies]
ndarray = { version = "0.14", features = ["blas"] }
blas-src = { version = "0.6.1", default-features = false, features = ["openblas"] }
openblas-src = { version = "0.9", default-features = false, features = ["cblas", "static"] }
ndarray = { version = "0.15", features = ["blas"] }
blas-src = { version = "0.8", default-features = false, features = ["openblas"] }
openblas-src = { version = "0.10", default-features = false, features = ["cblas", "static"] }
parking_lot = "0.11"
itertools = "0.10"
rocksdb = "0.15.0"
uuid = { version = "0.8", features = ["v4"] }
@@ -25,7 +25,6 @@ bincode = "1.3"
serde = { version = "~1.0", features = ["derive", "rc"] }
serde_json = "~1.0"
serde_cbor = "0.11.1"
rmp-serde = "~0.14"
ordered-float = "1.0"
thiserror = "1.0"
atomic_refcell = "0.1.6"
@@ -36,7 +35,15 @@ log = "0.4"
env_logger = "0.7.1"
geo = "0.17.0"
num-traits = "0.2.14"
rand = "0.8"
lru = "0.6.5"
bit-vec = "0.6"
fasthash = "0.4"
[[bench]]
name = "vector_search"
harness = false
harness = false
[[bench]]
name = "hnsw_build_graph"
harness = false

53
lib/segment/benches/hnsw_build_graph.rs Normal file
View File

@@ -0,0 +1,53 @@
mod prof;
use criterion::{Criterion, criterion_group, criterion_main};
use segment::types::{Distance, PointOffsetType};
use rand::thread_rng;
use segment::index::hnsw_index::graph_layers::GraphLayers;
use segment::index::hnsw_index::point_scorer::FilteredScorer;
use segment::fixtures::index_fixtures::{TestRawScorerProducer, FakeConditionChecker};
const NUM_VECTORS: usize = 10000;
const DIM: usize = 32;
const M: usize = 16;
const EF_CONSTRUCT: usize = 64;
const USE_HEURISTIC: bool = true;
fn hnsw_benchmark(c: &mut Criterion) {
let vector_holder = TestRawScorerProducer::new(DIM, NUM_VECTORS, Distance::Cosine);
let mut group = c.benchmark_group("hnsw-index-build-group");
group.sample_size(10);
group.bench_function("hnsw_index", |b| {
b.iter(|| {
let mut rng = thread_rng();
let mut graph_layers = GraphLayers::new(
NUM_VECTORS, M, M * 2, EF_CONSTRUCT, 10, USE_HEURISTIC,
);
let fake_condition_checker = FakeConditionChecker {};
for idx in 0..(NUM_VECTORS as PointOffsetType) {
let added_vector = vector_holder.vectors[idx as usize].to_vec();
let raw_scorer = vector_holder.get_raw_scorer(added_vector);
let scorer = FilteredScorer {
raw_scorer: &raw_scorer,
condition_checker: &fake_condition_checker,
filter: None,
};
let level = graph_layers.get_random_layer(&mut rng);
graph_layers.link_new_point(idx, level, &scorer);
}
})
});
group.finish();
}
criterion_group!{
name = benches;
config = Criterion::default().with_profiler(prof::FlamegraphProfiler::new(100));
targets = hnsw_benchmark
}
criterion_main!(benches);

79
lib/segment/benches/prof.rs Normal file
View File

@@ -0,0 +1,79 @@
use std::{fs::File, os::raw::c_int, path::Path};
use criterion::profiler::Profiler;
use pprof::ProfilerGuard;
use pprof::flamegraph::TextTruncateDirection;
/// Small custom profiler that can be used with Criterion to create a flamegraph for benchmarks.
/// Also see [the Criterion documentation on this][custom-profiler].
///
/// ## Example on how to enable the custom profiler:
///
/// ```
/// mod perf;
/// use perf::FlamegraphProfiler;
///
/// fn fibonacci_profiled(criterion: &mut Criterion) {
/// // Use the criterion struct as normal here.
/// }
///
/// fn custom() -> Criterion {
/// Criterion::default().with_profiler(FlamegraphProfiler::new())
/// }
///
/// criterion_group! {
/// name = benches;
/// config = custom();
/// targets = fibonacci_profiled
/// }
/// ```
///
/// The neat thing about this is that it will sample _only_ the benchmark, and not other stuff like
/// the setup process.
///
/// Further, it will only kick in if `--profile-time <time>` is passed to the benchmark binary.
/// A flamegraph will be created for each individual benchmark in its report directory under
/// `profile/flamegraph.svg`.
///
/// [custom-profiler]: https://bheisler.github.io/criterion.rs/book/user_guide/profiling.html#implementing-in-process-profiling-hooks
pub struct FlamegraphProfiler<'a> {
frequency: c_int,
active_profiler: Option<ProfilerGuard<'a>>,
}
impl<'a> FlamegraphProfiler<'a> {
#[allow(dead_code)]
pub fn new(frequency: c_int) -> Self {
FlamegraphProfiler {
frequency,
active_profiler: None,
}
}
}
impl<'a> Profiler for FlamegraphProfiler<'a> {
fn start_profiling(&mut self, _benchmark_id: &str, _benchmark_dir: &Path) {
self.active_profiler = Some(ProfilerGuard::new(self.frequency).unwrap());
}
fn stop_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) {
std::fs::create_dir_all(benchmark_dir).unwrap();
let flamegraph_path = benchmark_dir.join("flamegraph.svg");
eprintln!("\nflamegraph_path = {:#?}", flamegraph_path);
let flamegraph_file = File::create(&flamegraph_path)
.expect("File system error while creating flamegraph.svg");
let mut options = pprof::flamegraph::Options::default();
options.image_width = Some(2500);
options.text_truncate_direction = TextTruncateDirection::Left;
options.font_size = options.font_size / 3;
if let Some(profiler) = self.active_profiler.take() {
profiler
.report()
.build()
.unwrap()
.flamegraph_with_options(flamegraph_file, &mut options)
.expect("Error writing flamegraph");
}
}
}

23
lib/segment/benches/vector_search.rs
View File

@@ -4,27 +4,27 @@ use ndarray::{Array, Array1, Array2, ArrayBase, ShapeBuilder, Axis};
use tempdir::TempDir;
use segment::spaces::tools::{peek_top_scores, peek_top_scores_iterable};
use segment::types::Distance;
use segment::types::{Distance, VectorElementType, PointOffsetType};
use segment::vector_storage::simple_vector_storage::SimpleVectorStorage;
use segment::vector_storage::vector_storage::{ScoredPointOffset, VectorStorage};
const NUM_VECTORS: usize = 50000;
const DIM: usize = 1000; // Larger dimensionality - greater the BLAS advantage
fn random_vector(size: usize) -> Vec<f32> {
let mut vec: Vec<f32> = Vec::with_capacity(size);
fn random_vector(size: usize) -> Vec<VectorElementType> {
let mut vec: Vec<VectorElementType> = Vec::with_capacity(size);
for _ in 0..vec.capacity() {
vec.push(rand::random());
};
return vec;
}
fn init_vector_storage(dir: &TempDir, dim: usize, num: usize) -> SimpleVectorStorage {
let mut storage = SimpleVectorStorage::open(dir.path(), dim).unwrap();
fn init_vector_storage(dir: &TempDir, dim: usize, num: usize, dist: Distance) -> SimpleVectorStorage {
let mut storage = SimpleVectorStorage::open(dir.path(), dim, dist).unwrap();
for _i in 0..num {
let vector: Vec<f32> = random_vector(dim);
storage.put_vector(&vector).unwrap();
let vector: Vec<VectorElementType> = random_vector(dim);
storage.put_vector(vector).unwrap();
}
storage
@@ -34,12 +34,12 @@ fn benchmark_naive(c: &mut Criterion) {
let dir = TempDir::new("storage_dir").unwrap();
let dist = Distance::Dot;
let storage = init_vector_storage(&dir, DIM, NUM_VECTORS);
let storage = init_vector_storage(&dir, DIM, NUM_VECTORS, dist);
c.bench_function("storage vector search",
|b| b.iter(|| {
let vector = random_vector(DIM);
storage.score_all(&vector, 10, &dist)
storage.score_all(&vector, 10)
}));
}
@@ -64,9 +64,8 @@ fn benchmark_ndarray(c: &mut Criterion) {
.cloned()
.enumerate()
.map(
|(idx, score)| ScoredPointOffset { idx, score }),
10,
&Distance::Dot,
|(idx, score)| ScoredPointOffset { idx: idx as PointOffsetType, score }),
10
);
}));
}

24
lib/segment/src/common/file_operations.rs
View File

@@ -1,5 +1,5 @@
use std::fs::File;
use std::io::{Read, Write};
use std::io::{Read, Write, BufWriter};
use crate::entry::entry_point::{OperationError, OperationResult};
use serde::Serialize;
use std::path::Path;
@@ -7,6 +7,16 @@ use atomicwrites::AtomicFile;
use atomicwrites::OverwriteBehavior::AllowOverwrite;
use serde::de::DeserializeOwned;
pub fn atomic_save_bin<N: DeserializeOwned + Serialize>(path: &Path, object: &N) -> OperationResult<()> {
let af = AtomicFile::new(path, AllowOverwrite);
af.write(|f| {
let mut writer = BufWriter::new(f);
bincode::serialize_into(&mut writer, object)
})?;
Ok(())
}
pub fn atomic_save_json<N: DeserializeOwned + Serialize>(path: &Path, object: &N) -> OperationResult<()> {
let af = AtomicFile::new(path, AllowOverwrite);
let state_bytes = serde_json::to_vec(object).unwrap();
@@ -28,5 +38,17 @@ pub fn read_json<N: DeserializeOwned + Serialize>(path: &Path) -> OperationResul
})
})?;
Ok(result)
}
pub fn read_bin<N: DeserializeOwned + Serialize>(path: &Path) -> OperationResult<N> {
let mut file = File::open(path)?;
let result: N = bincode::deserialize_from(&mut file).or_else(|err| {
Err(OperationError::ServiceError {
description: format!("Failed to read data {}. Error: {}", path.to_str().unwrap(), err)
})
})?;
Ok(result)
}

3
lib/segment/src/common/mod.rs
View File

@@ -1,2 +1,3 @@
pub mod file_operations;
pub mod error_logging;
pub mod error_logging;
pub mod utils;

6
lib/segment/src/common/utils.rs Normal file
View File

@@ -0,0 +1,6 @@
pub fn rev_range(a: usize, b: usize) -> impl Iterator<Item=usize> {
(b + 1..=a).rev()
}

56
lib/segment/src/fixtures/index_fixtures.rs Normal file
View File

@@ -0,0 +1,56 @@
use ndarray::{Array1, Array};
use crate::types::{VectorElementType, PointOffsetType, Distance, Filter};
use crate::spaces::metric::Metric;
use crate::spaces::tools::mertic_object;
use crate::vector_storage::simple_vector_storage::SimpleRawScorer;
use crate::payload_storage::payload_storage::ConditionChecker;
use itertools::Itertools;
use rand::prelude::ThreadRng;
use rand::Rng;
use bit_vec::BitVec;
pub fn random_vector(rnd_gen: &mut ThreadRng, size: usize) -> Vec<VectorElementType> {
(0..size).map(|_| rnd_gen.gen()).collect()
}
pub struct FakeConditionChecker {}
impl ConditionChecker for FakeConditionChecker {
fn check(&self, _point_id: PointOffsetType, _query: &Filter) -> bool { true }
}
pub struct TestRawScorerProducer {
pub vectors: Vec<Array1<VectorElementType>>,
pub deleted: BitVec,
pub metric: Box<dyn Metric>,
}
impl TestRawScorerProducer {
pub fn new(dim: usize, num_vectors: usize, distance: Distance) -> Self {
let mut rnd = rand::thread_rng();
let metric = mertic_object(&distance);
let vectors = (0..num_vectors)
.map(|_x| metric.preprocess(random_vector(&mut rnd, dim)))
.collect_vec();
TestRawScorerProducer {
vectors: vectors.into_iter().map(|v| Array::from(v)).collect(),
deleted: BitVec::from_elem(num_vectors, false),
metric,
}
}
pub fn get_raw_scorer(&self, query: Vec<VectorElementType>) -> SimpleRawScorer {
SimpleRawScorer {
query: Array::from(self.metric.preprocess(query)),
metric: &self.metric,
vectors: &self.vectors,
deleted: &self.deleted,
}
}
}

2
lib/segment/src/fixtures/mod.rs Normal file
View File

@@ -0,0 +1,2 @@
pub mod index_fixtures;
pub mod payload_fixtures;

114
lib/segment/src/fixtures/payload_fixtures.rs Normal file
View File

@@ -0,0 +1,114 @@
use rand::prelude::ThreadRng;
use rand::seq::SliceRandom;
use crate::types::{Filter, PayloadType, VectorElementType, Condition, FieldCondition, Match, Range as RangeCondition};
use rand::Rng;
use itertools::Itertools;
use std::ops::Range;
const ADJECTIVE: &'static [&'static str] = &[
"jobless",
"rightful",
"breakable",
"impartial",
"shocking",
"faded",
"phobic",
"overt",
"like",
"wide-eyed",
"broad",
];
const NOUN: &'static [&'static str] = &[
"territory",
"jam",
"neck",
"chicken",
"cap",
"kiss",
"veil",
"trail",
"size",
"digestion",
"rod",
"seed",
];
const INT_RANGE: Range<i64> = 0..500;
pub fn random_keyword(rnd_gen: &mut ThreadRng) -> String {
let random_adj = ADJECTIVE.choose(rnd_gen).unwrap();
let random_noun = NOUN.choose(rnd_gen).unwrap();
format!("{} {}", random_adj, random_noun)
}
pub fn random_keyword_payload(rnd_gen: &mut ThreadRng) -> PayloadType {
PayloadType::Keyword(vec![random_keyword(rnd_gen)])
}
pub fn random_int_payload(rnd_gen: &mut ThreadRng, num_values: usize) -> PayloadType {
PayloadType::Integer((0..num_values).map(|_| rnd_gen.gen_range(INT_RANGE)).collect_vec())
}
pub fn random_vector(rnd_gen: &mut ThreadRng, size: usize) -> Vec<VectorElementType> {
(0..size).map(|_| rnd_gen.gen()).collect()
}
pub fn random_field_condition(rnd_gen: &mut ThreadRng) -> Condition {
let kv_or_int: bool = rnd_gen.gen();
match kv_or_int {
true => Condition::Field(FieldCondition {
key: "kvd".to_string(),
r#match: Some(Match {
keyword: Some(random_keyword(rnd_gen)),
integer: None,
}),
range: None,
geo_bounding_box: None,
geo_radius: None,
}),
false => Condition::Field(FieldCondition {
key: "int".to_string(),
r#match: None,
range: Some(RangeCondition {
lt: None,
gt: None,
gte: Some(rnd_gen.gen_range(INT_RANGE) as f64),
lte: Some(rnd_gen.gen_range(INT_RANGE) as f64),
}),
geo_bounding_box: None,
geo_radius: None,
})
}
}
pub fn random_filter(rnd_gen: &mut ThreadRng) -> Filter {
let mut rnd1 = rand::thread_rng();
let should_conditions = (0..=2)
.take_while(|_| rnd1.gen::<f64>() > 0.6)
.map(|_| random_field_condition(rnd_gen))
.collect_vec();
let should_conditions_opt = match should_conditions.is_empty() {
false => Some(should_conditions),
true => None,
};
let must_conditions = (0..=2)
.take_while(|_| rnd1.gen::<f64>() > 0.6)
.map(|_| random_field_condition(rnd_gen))
.collect_vec();
let must_conditions_opt = match must_conditions.is_empty() {
false => Some(must_conditions),
true => None,
};
Filter {
should: should_conditions_opt,
must: must_conditions_opt,
must_not: None,
}
}

13
lib/segment/src/index/field_index/field_index.rs
View File

@@ -1,14 +1,19 @@
use serde::{Deserialize, Serialize};
use crate::index::field_index::CardinalityEstimation;
use crate::index::field_index::{CardinalityEstimation, PayloadBlockCondition};
use crate::index::field_index::map_index::PersistedMapIndex;
use crate::index::field_index::numeric_index::PersistedNumericIndex;
use crate::types::{FieldCondition, FloatPayloadType, IntPayloadType, PayloadType, PointOffsetType};
use crate::types::{FieldCondition, FloatPayloadType, IntPayloadType, PayloadType, PointOffsetType, PayloadKeyType};
pub trait PayloadFieldIndex {
/// Get iterator over points fitting given `condition`
fn filter(&self, condition: &FieldCondition) -> Option<Box<dyn Iterator<Item=PointOffsetType> + '_>>;
/// Return estimation of points amount which satisfy given condition
fn estimate_cardinality(&self, condition: &FieldCondition) -> Option<CardinalityEstimation>;
/// Iterate conditions for payload blocks with minimum size of `threshold`
/// Required for building HNSW index
fn payload_blocks(&self, threshold: usize, key: PayloadKeyType) -> Box<dyn Iterator<Item=PayloadBlockCondition> + '_>;
}
pub trait PayloadFieldIndexBuilder {
@@ -46,4 +51,8 @@ impl PayloadFieldIndex for FieldIndex {
fn estimate_cardinality(&self, condition: &FieldCondition) -> Option<CardinalityEstimation> {
self.get_payload_field_index().estimate_cardinality(condition)
}
fn payload_blocks(&self, threshold: usize, key: PayloadKeyType) -> Box<dyn Iterator<Item=PayloadBlockCondition> + '_> {
self.get_payload_field_index().payload_blocks(threshold, key)
}
}

52
lib/segment/src/index/field_index/map_index.rs
View File

@@ -4,9 +4,9 @@ use std::{mem, iter};
use serde::{Deserialize, Serialize};
use crate::index::field_index::{CardinalityEstimation, PrimaryCondition};
use crate::index::field_index::{CardinalityEstimation, PrimaryCondition, PayloadBlockCondition};
use crate::index::field_index::field_index::{FieldIndex, PayloadFieldIndex, PayloadFieldIndexBuilder};
use crate::types::{IntPayloadType, PayloadType, PointOffsetType, FieldCondition};
use crate::types::{IntPayloadType, PayloadType, PointOffsetType, FieldCondition, PayloadKeyType, Match};
#[derive(Serialize, Deserialize)]
pub struct PersistedMapIndex<N: Hash + Eq + Clone> {
@@ -75,6 +75,28 @@ impl PayloadFieldIndex for PersistedMapIndex<String> {
})
)
}
fn payload_blocks(&self, threshold: usize, key: PayloadKeyType) -> Box<dyn Iterator<Item=PayloadBlockCondition> + '_> {
let iter = self.map
.iter()
.filter(move |(_value, point_ids)| point_ids.len() > threshold)
.map(move |(value, point_ids)| {
PayloadBlockCondition {
condition: FieldCondition {
key: key.clone(),
r#match: Some(Match {
keyword: Some(value.to_owned()),
integer: None,
}),
range: None,
geo_bounding_box: None,
geo_radius: None,
},
cardinality: point_ids.len(),
}
});
Box::new(iter)
}
}
impl PayloadFieldIndex for PersistedMapIndex<IntPayloadType> {
@@ -94,10 +116,32 @@ impl PayloadFieldIndex for PersistedMapIndex<IntPayloadType> {
estimation
}))
}
fn payload_blocks(&self, threshold: usize, key: PayloadKeyType) -> Box<dyn Iterator<Item=PayloadBlockCondition> + '_> {
let iter = self.map
.iter()
.filter(move |(_value, point_ids)| point_ids.len() >= threshold)
.map(move |(value, point_ids)| {
PayloadBlockCondition {
condition: FieldCondition {
key: key.clone(),
r#match: Some(Match {
keyword: None,
integer: Some(*value),
}),
range: None,
geo_bounding_box: None,
geo_radius: None,
},
cardinality: point_ids.len(),
}
});
Box::new(iter)
}
}
impl PayloadFieldIndexBuilder for PersistedMapIndex<String> {
fn add(&mut self, id: usize, value: &PayloadType) {
fn add(&mut self, id: PointOffsetType, value: &PayloadType) {
match value {
PayloadType::Keyword(keywords) => self.add_many(id, keywords),
_ => panic!("Unexpected payload type: {:?}", value)
@@ -114,7 +158,7 @@ impl PayloadFieldIndexBuilder for PersistedMapIndex<String> {
}
impl PayloadFieldIndexBuilder for PersistedMapIndex<IntPayloadType> {
fn add(&mut self, id: usize, value: &PayloadType) {
fn add(&mut self, id: PointOffsetType, value: &PayloadType) {
match value {
PayloadType::Integer(numbers) => self.add_many(id, numbers),
_ => panic!("Unexpected payload type: {:?}", value)

6
lib/segment/src/index/field_index/mod.rs
View File

@@ -13,6 +13,12 @@ pub enum PrimaryCondition {
Ids(HashSet<PointOffsetType>),
}
#[derive(Debug, Clone)]
pub struct PayloadBlockCondition {
pub condition: FieldCondition,
pub cardinality: usize
}
#[derive(Debug)]
pub struct CardinalityEstimation {
/// Conditions that could be used to mane a primary point selection.

40
lib/segment/src/index/field_index/numeric_index.rs
View File

@@ -6,9 +6,9 @@ use num_traits::ToPrimitive;
use ordered_float::OrderedFloat;
use serde::{Deserialize, Serialize};
use crate::index::field_index::{CardinalityEstimation, PrimaryCondition};
use crate::index::field_index::{CardinalityEstimation, PrimaryCondition, PayloadBlockCondition};
use crate::index::field_index::field_index::{FieldIndex, PayloadFieldIndex, PayloadFieldIndexBuilder};
use crate::types::{FloatPayloadType, IntPayloadType, PayloadType, PointOffsetType, Range, FieldCondition};
use crate::types::{FloatPayloadType, IntPayloadType, PayloadType, PointOffsetType, Range, FieldCondition, PayloadKeyType};
#[derive(Debug, Deserialize, Serialize, Clone)]
pub struct Element<N> {
@@ -68,12 +68,14 @@ impl<N: ToPrimitive + Clone> PersistedNumericIndex<N> {
(0, 0)
} else {
(lower_index, upper_index)
}
};
}
pub fn range_cardinality(&self, range: &Range) -> CardinalityEstimation {
let (lower_index, upper_index) = self.search_range(range);
// ToDo: Check if there is a more precise implementation for multiple values
let values_count: i64 = upper_index as i64 - lower_index as i64;
let total_values = self.elements.len() as i64;
let value_per_point = total_values as f64 / self.points_count as f64;
@@ -125,6 +127,38 @@ impl<N: ToPrimitive + Clone> PayloadFieldIndex for PersistedNumericIndex<N> {
cardinality
})
}
fn payload_blocks(&self, threshold: usize, key: PayloadKeyType) -> Box<dyn Iterator<Item=PayloadBlockCondition> + '_> {
// Creates half-overlapped ranges of points.
let num_elements = self.elements.len();
let value_per_point = num_elements as f64 / self.points_count as f64;
let effective_threshold = (threshold as f64 * value_per_point) as usize;
let iter = (0..num_elements).step_by(effective_threshold / 2).map(move |init_offset| {
let upper_index = min(num_elements - 1, init_offset + effective_threshold);
let upper_value = self.elements[upper_index].value.to_f64();
let lower_value = self.elements[init_offset].value.to_f64();
PayloadBlockCondition {
condition: FieldCondition {
key: key.clone(),
r#match: None,
range: Some(Range {
lt: None,
gt: None,
gte: lower_value,
lte: upper_value,
}),
geo_bounding_box: None,
geo_radius: None,
},
cardinality: ((upper_index - init_offset) as f64 / value_per_point) as usize,
}
});
Box::new(iter)
}
}
impl PayloadFieldIndexBuilder for PersistedNumericIndex<FloatPayloadType> {

105
lib/segment/src/index/hnsw_index/build_cache.rs Normal file
View File

@@ -0,0 +1,105 @@
use crate::types::{ScoreType, PointOffsetType};
use std::hash::{Hasher, Hash};
use std::cmp::{min, max};
use fasthash::SeaHasher;
#[derive(Hash, Eq, PartialEq, Clone, Debug)]
struct PointPair {
a: PointOffsetType,
b: PointOffsetType,
}
impl PointPair {
pub fn new(a: PointOffsetType, b: PointOffsetType) -> Self {
PointPair {
a: min(a, b),
b: max(a, b),
}
}
}
#[derive(Clone, Debug)]
struct CacheObj {
points: PointPair,
value: ScoreType,
}
#[derive(Debug)]
pub struct DistanceCache {
cache: Vec<Option<CacheObj>>,
pub hits: usize,
pub misses: usize
}
#[allow(dead_code)]
impl DistanceCache {
fn hasher() -> impl Hasher {
SeaHasher::new()
}
pub fn new(size: usize) -> Self {
let mut cache = Vec::with_capacity(size);
cache.resize(size, None);
DistanceCache {
cache,
hits: 0,
misses: 0
}
}
pub fn get(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> Option<ScoreType> {
let points = PointPair::new(point_a, point_b);
let mut s = DistanceCache::hasher();
points.hash(&mut s);
let idx = s.finish() as usize % self.cache.len();
self.cache[idx].as_ref().and_then(|x| if x.points == points { Some(x.value) } else { None })
}
pub fn put(&mut self, point_a: PointOffsetType, point_b: PointOffsetType, value: ScoreType) {
let points = PointPair::new(point_a, point_b);
let mut s = DistanceCache::hasher();
points.hash(&mut s);
let idx = s.finish() as usize % self.cache.len();
self.cache[idx] = Some(CacheObj { points, value })
}
}
impl Default for DistanceCache {
fn default() -> Self {
DistanceCache::new(0)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_cache() {
let mut cache = DistanceCache::new(1000);
cache.put(100, 10, 0.8);
cache.put(10, 101, 0.7);
cache.put(10, 110, 0.1);
assert_eq!(cache.get(12, 99), None);
assert_eq!(cache.get(10, 100), Some(0.8));
assert_eq!(cache.get(10, 101), Some(0.7));
}
#[test]
fn test_collision() {
let mut cache = DistanceCache::new(1);
cache.put(1, 2, 0.8);
cache.put(3, 4, 0.7);
assert_eq!(cache.get(1, 2), None);
assert_eq!(cache.get(2, 1), None);
assert_eq!(cache.get(4, 3), Some(0.7));
}
}

22
lib/segment/src/index/hnsw_index/build_condition_checker.rs Normal file
View File

@@ -0,0 +1,22 @@
use crate::payload_storage::payload_storage::ConditionChecker;
use crate::types::{Filter, PointOffsetType};
use crate::index::visited_pool::VisitedList;
pub struct BuildConditionChecker {
pub filter_list: VisitedList
}
impl BuildConditionChecker {
pub fn new(list_size: usize) -> Self {
BuildConditionChecker {
filter_list: VisitedList::new(list_size)
}
}
}
impl ConditionChecker for BuildConditionChecker {
fn check(&self, point_id: PointOffsetType, _query: &Filter) -> bool {
self.filter_list.check(point_id)
}
}

45
lib/segment/src/index/hnsw_index/config.rs Normal file
View File

@@ -0,0 +1,45 @@
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use crate::entry::entry_point::OperationResult;
use crate::common::file_operations::{read_json, atomic_save_json};
pub const HNSW_INDEX_CONFIG_FILE: &str = "hnsw_config.json";
#[derive(Debug, Deserialize, Serialize, Copy, Clone, PartialEq)]
pub struct HnswGraphConfig {
pub m: usize,
/// Requested M
pub m0: usize,
/// Actual M on level 0
pub ef_construct: usize,
/// Number of neighbours to search on construction
pub ef: usize,
/// Minimal number of vectors to perform indexing
pub indexing_threshold: usize,
}
impl HnswGraphConfig {
pub fn new(m: usize, ef_construct: usize, indexing_threshold: usize) -> Self {
HnswGraphConfig {
m,
m0: m * 2,
ef_construct,
ef: ef_construct,
indexing_threshold
}
}
pub fn get_config_path(path: &Path) -> PathBuf {
path.join(HNSW_INDEX_CONFIG_FILE)
}
pub fn load(path: &Path) -> OperationResult<Self> {
read_json(path)
}
pub fn save(&self, path: &Path) -> OperationResult<()> {
atomic_save_json(path, self)
}
}

135
lib/segment/src/index/hnsw_index/entry_points.rs Normal file
View File

@@ -0,0 +1,135 @@
use serde::{Deserialize, Serialize};
use crate::types::PointOffsetType;
use std::cmp::Ordering;
use crate::spaces::tools::FixedLengthPriorityQueue;
#[derive(Deserialize, Serialize, Clone, Debug, PartialEq)]
pub struct EntryPoint {
pub point_id: PointOffsetType,
pub level: usize,
}
impl Eq for EntryPoint {}
impl PartialOrd for EntryPoint {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for EntryPoint {
fn cmp(&self, other: &Self) -> Ordering {
self.level.cmp(&other.level)
}
}
#[derive(Deserialize, Serialize, Clone, Debug)]
pub struct EntryPoints {
entry_points: Vec<EntryPoint>,
extra_entry_points: FixedLengthPriorityQueue<EntryPoint>,
}
impl EntryPoints {
pub fn new(extra_entry_points: usize) -> Self {
EntryPoints {
entry_points: vec![],
extra_entry_points: FixedLengthPriorityQueue::new(extra_entry_points),
}
}
pub fn merge_from_other(&mut self, mut other: EntryPoints) {
self.entry_points.append(&mut other.entry_points);
// Do not merge `extra_entry_points` to prevent duplications
}
pub fn new_point<F>(&mut self, new_point: PointOffsetType, level: usize, checker: F) -> Option<EntryPoint>
where F: Fn(PointOffsetType) -> bool
{
// there are 3 cases:
// - There is proper entry point for a new point higher or same level - return the point
// - The new point is higher than any alternative - return the next best thing
// - There is no point and alternatives - return None
for i in 0..self.entry_points.len() {
let candidate = &self.entry_points[i];
if !checker(candidate.point_id) {
continue; // Checkpoint does not fulfil filtering conditions. Hence, does not "exists"
}
// Found checkpoint candidate
return if candidate.level >= level {
// The good checkpoint exists.
// Return it, and also try to save given if required
self.extra_entry_points.push(EntryPoint {
point_id: new_point,
level,
});
Some(candidate.clone())
} else {
// The current point is better than existing
let entry = self.entry_points[i].clone();
self.entry_points[i] = EntryPoint {
point_id: new_point,
level,
};
self.extra_entry_points.push(entry.clone());
Some(entry)
};
}
// No entry points found. Create a new one and return self
let new_entry = EntryPoint {
point_id: new_point,
level,
};
self.entry_points.push(new_entry.clone());
None
}
/// Find the highest EntryPoint which satisfies filtering condition of `checker`
pub fn get_entry_point<F>(&self, checker: F) -> Option<EntryPoint>
where F: Fn(PointOffsetType) -> bool
{
self.entry_points.iter()
.filter(|entry| checker(entry.point_id))
.cloned().next().or_else(|| {
// Searching for at least some entry point
self.extra_entry_points
.iter()
.filter(|entry| checker(entry.point_id))
.cloned()
.max_by_key(|ep| ep.level)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use rand::Rng;
#[test]
fn test_entry_points() {
let mut points = EntryPoints::new(10);
let mut rnd = rand::thread_rng();
for i in 0..1000 {
let level = rnd.gen_range(0..10000);
points.new_point(i, level, |_x| true);
}
assert_eq!(points.entry_points.len(), 1);
assert_eq!(points.extra_entry_points.len(), 10);
assert!(points.entry_points[0].level > 1);
for i in 1000..2000 {
let level = rnd.gen_range(0..10000);
points.new_point(i, level, |x| x % 5 == i % 5);
}
assert_eq!(points.entry_points.len(), 5);
assert_eq!(points.extra_entry_points.len(), 10);
}
}

665
lib/segment/src/index/hnsw_index/graph_layers.rs Normal file
View File

@@ -0,0 +1,665 @@
use serde::{Deserialize, Serialize};
use crate::types::{PointOffsetType, ScoreType};
use crate::spaces::tools::FixedLengthPriorityQueue;
use std::cmp::{max, min};
use std::path::{Path, PathBuf};
use crate::entry::entry_point::OperationResult;
use crate::common::file_operations::{read_bin, atomic_save_bin};
use crate::index::hnsw_index::point_scorer::FilteredScorer;
use crate::index::hnsw_index::entry_points::EntryPoints;
use crate::vector_storage::vector_storage::ScoredPointOffset;
use crate::index::visited_pool::{VisitedList, VisitedPool};
use crate::index::hnsw_index::search_context::SearchContext;
use crate::common::utils::rev_range;
use rand::distributions::Uniform;
use rand::prelude::ThreadRng;
use rand::Rng;
use std::collections::BinaryHeap;
use itertools::Itertools;
pub type LinkContainer = Vec<PointOffsetType>;
pub type LayersContainer = Vec<LinkContainer>;
pub const HNSW_GRAPH_FILE: &str = "graph.bin";
#[derive(Deserialize, Serialize, Debug)]
pub struct GraphLayers {
max_level: usize,
m: usize,
m0: usize,
ef_construct: usize,
level_factor: f64,
// Exclude points according to "not closer than base" heuristic?
use_heuristic: bool,
// Factor of level probability
links_layers: Vec<LayersContainer>,
entry_points: EntryPoints,
// Fields used on construction phase only
#[serde(skip)]
visited_pool: VisitedPool,
}
/// Object contains links between nodes for HNSW search
///
/// Assume all scores are similarities. Larger score = closer points
impl GraphLayers {
pub fn new_with_params(
num_vectors: usize, // Initial number of points in index
m: usize, // Expected M for non-first layer
m0: usize, // Expected M for first layer
ef_construct: usize,
entry_points_num: usize, // Depends on number of points
use_heuristic: bool,
reserve: bool
) -> Self {
let mut links_layers: Vec<LayersContainer> = vec![];
for _i in 0..num_vectors {
let mut links: LinkContainer = Vec::new();
if reserve {
links.reserve(m0);
}
links_layers.push(vec![links]);
}
GraphLayers {
max_level: 0,
m,
m0,
ef_construct,
level_factor: 1.0 / (m as f64).ln(),
use_heuristic,
links_layers,
entry_points: EntryPoints::new(entry_points_num),
visited_pool: VisitedPool::new(),
}
}
pub fn new(
num_vectors: usize, // Initial number of points in index
m: usize, // Expected M for non-first layer
m0: usize, // Expected M for first layer
ef_construct: usize,
entry_points_num: usize, // Depends on number of points
use_heuristic: bool,
) -> Self {
Self::new_with_params(num_vectors, m, m0, ef_construct, entry_points_num, use_heuristic, true)
}
fn num_points(&self) -> usize { self.links_layers.len() }
pub fn point_level(&self, point_id: PointOffsetType) -> usize {
self.links_layers[point_id as usize].len() - 1
}
/// Get links of current point
fn links(&self, point_id: PointOffsetType, level: usize) -> &LinkContainer {
&self.links_layers[point_id as usize][level]
}
/// Get M based on current level
fn get_m(&self, level: usize) -> usize {
return if level == 0 { self.m0 } else { self.m };
}
/// Generate random level for a new point, according to geometric distribution
pub fn get_random_layer(&self, thread_rng: &mut ThreadRng) -> usize {
let distribution = Uniform::new(0.0, 1.0);
let sample: f64 = thread_rng.sample(distribution);
let picked_level = -sample.ln() * self.level_factor;
return picked_level.round() as usize;
}
fn set_levels(&mut self, point_id: PointOffsetType, level: usize) {
if self.links_layers.len() <= point_id as usize {
self.links_layers.resize(point_id as usize, vec![]);
}
let point_layers = &mut self.links_layers[point_id as usize];
while point_layers.len() <= level {
let mut links = vec![];
links.reserve(self.m);
point_layers.push(links)
}
self.max_level = max(level, self.max_level);
}
/// Greedy search for closest points within a single graph layer
fn _search_on_level(&self, searcher: &mut SearchContext, level: usize, visited_list: &mut VisitedList, points_scorer: &FilteredScorer) {
while let Some(index) = searcher.candidates.pop() {
let mut links_iter = self.links(index, level)
.iter()
.cloned()
.filter(|point_id| !visited_list.check_and_update_visited(*point_id));
points_scorer.score_iterable_points(
&mut links_iter,
self.get_m(level),
|score_point| searcher.process_candidate(score_point),
);
}
}
fn search_on_level(&self, level_entry: ScoredPointOffset, level: usize, ef: usize, points_scorer: &FilteredScorer) -> FixedLengthPriorityQueue<ScoredPointOffset> {
let mut visited_list = self.visited_pool.get(self.num_points());
visited_list.check_and_update_visited(level_entry.idx);
let mut search_context = SearchContext::new(level_entry, ef);
self._search_on_level(&mut search_context, level, &mut visited_list, points_scorer);
self.visited_pool.return_back(visited_list);
search_context.nearest
}
/// Greedy searches for entry point of level `target_level`.
/// Beam size is 1.
fn search_entry(&self, entry_point: PointOffsetType, top_level: usize, target_level: usize, points_scorer: &FilteredScorer) -> ScoredPointOffset {
let mut current_point = ScoredPointOffset {
idx: entry_point,
score: points_scorer.score_point(entry_point),
};
for level in rev_range(top_level, target_level) {
let mut changed = true;
while changed {
changed = false;
let mut links = self.links(current_point.idx, level).iter().cloned();
points_scorer.score_iterable_points(
&mut links,
self.get_m(level),
|score_point| {
if score_point.score > current_point.score {
changed = true;
current_point = score_point;
}
},
);
}
}
current_point
}
/// Connect new point to links, so that links contains only closest points
fn connect_new_point<F>(
links: &mut LinkContainer,
new_point_id: PointOffsetType,
target_point_id: PointOffsetType,
level_m: usize,
mut score_internal: F,
)
where F: FnMut(PointOffsetType, PointOffsetType) -> ScoreType
{
// ToDo: binary search here ? (most likely does not worth it)
let new_to_target = score_internal(target_point_id, new_point_id);
let mut id_to_insert = links.len();
for i in 0..links.len() {
let target_to_link = score_internal(target_point_id, links[i]);
if target_to_link < new_to_target {
id_to_insert = i;
break;
}
}
if links.len() < level_m {
links.insert(id_to_insert, new_point_id)
} else {
if id_to_insert != links.len() {
links.pop();
links.insert(id_to_insert, new_point_id)
}
}
}
/// https://github.com/nmslib/hnswlib/issues/99
fn select_candidate_with_heuristic_from_sorted<F>(
candidates: impl Iterator<Item=ScoredPointOffset>,
m: usize,
mut score_internal: F,
) -> Vec<PointOffsetType>
where F: FnMut(PointOffsetType, PointOffsetType) -> ScoreType
{
let mut result_list = vec![];
result_list.reserve(m);
for current_closest in candidates {
if result_list.len() >= m { break; }
let mut is_good = true;
for selected_point in result_list.iter().cloned() {
let dist_to_already_selected = score_internal(current_closest.idx, selected_point);
if dist_to_already_selected > current_closest.score {
is_good = false;
break;
}
}
if is_good { result_list.push(current_closest.idx) }
}
result_list
}
/// https://github.com/nmslib/hnswlib/issues/99
fn select_candidates_with_heuristic<F>(
candidates: FixedLengthPriorityQueue<ScoredPointOffset>,
m: usize,
score_internal: F,
) -> Vec<PointOffsetType>
where F: FnMut(PointOffsetType, PointOffsetType) -> ScoreType {
let closest_iter = candidates.into_iter();
return Self::select_candidate_with_heuristic_from_sorted(closest_iter, m, score_internal);
}
pub fn link_new_point(&mut self, point_id: PointOffsetType, level: usize, points_scorer: &FilteredScorer) {
// Check if there is an suitable entry point
// - entry point level if higher or equal
// - it satisfies filters
self.set_levels(point_id, level);
let entry_point_opt = self.entry_points.new_point(
point_id,
level,
|point_id| points_scorer.check_point(point_id),
);
match entry_point_opt {
// New point is a new empty entry (for this filter, at least)
// We can't do much here, so just quit
None => {}
// Entry point found.
Some(entry_point) => {
let mut level_entry = if entry_point.level > level {
// The entry point is higher than a new point
// Let's find closest one on same level
// greedy search for a single closest point
self.search_entry(
entry_point.point_id,
entry_point.level,
level,
points_scorer,
)
} else {
ScoredPointOffset {
idx: entry_point.point_id,
score: points_scorer.score_internal(point_id, entry_point.point_id),
}
};
// minimal common level for entry points
let linking_level = min(level, entry_point.level);
let scorer = |a, b| points_scorer.score_internal(a, b);
for curr_level in (0..=linking_level).rev() {
let level_m = self.get_m(curr_level);
let nearest_points = self.search_on_level(
level_entry, curr_level, self.ef_construct, points_scorer,
);
if self.use_heuristic {
let selected_nearest = Self::select_candidates_with_heuristic(
nearest_points, level_m, scorer);
self.links_layers[point_id as usize][curr_level].clone_from(&selected_nearest);
for other_point in selected_nearest.iter().cloned() {
let other_point_links = &mut self.links_layers[other_point as usize][curr_level];
if other_point_links.len() < level_m {
// If linked point is lack of neighbours
other_point_links.push(point_id);
} else {
let mut candidates = BinaryHeap::with_capacity(level_m + 1);
candidates.push(ScoredPointOffset {
idx: point_id,
score: scorer(point_id, other_point),
});
for other_point_link in other_point_links.iter().take(level_m).cloned() {
candidates.push(ScoredPointOffset {
idx: other_point_link,
score: scorer(other_point_link, other_point),
});
}
let selected_candidates = Self::select_candidate_with_heuristic_from_sorted(
candidates.into_sorted_vec().into_iter().rev(),
level_m,
scorer,
);
for (idx, selected) in selected_candidates.iter().cloned().enumerate() {
other_point_links[idx] = selected;
}
}
}
} else {
for nearest_point in nearest_points.iter() {
Self::connect_new_point(
&mut self.links_layers[point_id as usize][curr_level],
nearest_point.idx,
point_id,
level_m,
scorer,
);
Self::connect_new_point(
&mut self.links_layers[nearest_point.idx as usize][curr_level],
point_id,
nearest_point.idx,
level_m,
scorer,
);
if nearest_point.score > level_entry.score {
level_entry = nearest_point.clone()
}
}
}
}
}
}
}
pub fn merge_from_other(&mut self, other: GraphLayers) {
let mut visited_list = self.visited_pool.get(self.num_points());
if other.links_layers.len() > self.links_layers.len() {
self.links_layers.resize(other.links_layers.len(), vec![])
}
for (point_id, layers) in other.links_layers.into_iter().enumerate() {
let current_layers = &mut self.links_layers[point_id];
for (level, other_links) in layers.into_iter().enumerate() {
if current_layers.len() <= level {
current_layers.push(other_links)
} else {
visited_list.next_iteration();
let current_links = &mut current_layers[level];
current_links.iter().cloned().for_each(|x| {visited_list.check_and_update_visited(x);});
for other_link in other_links.into_iter().filter(|x| !visited_list.check_and_update_visited(*x)) {
current_links.push(other_link)
}
}
}
}
self.entry_points.merge_from_other(other.entry_points);
self.visited_pool.return_back(visited_list);
}
pub fn search(&self, top: usize, ef: usize, points_scorer: &FilteredScorer) -> Vec<ScoredPointOffset> {
let entry_point = match self.entry_points.get_entry_point(|point_id| points_scorer.check_point(point_id)) {
None => return vec![],
Some(ep) => ep
};
let zero_level_entry = self.search_entry(
entry_point.point_id,
entry_point.level,
0,
points_scorer,
);
let nearest = self.search_on_level(zero_level_entry, 0, max(top, ef), points_scorer);
nearest.into_iter().take(top).collect_vec()
}
pub fn get_path(path: &Path) -> PathBuf {
path.join(HNSW_GRAPH_FILE)
}
pub fn load(path: &Path) -> OperationResult<Self> {
read_bin(path)
}
pub fn save(&self, path: &Path) -> OperationResult<()> {
atomic_save_bin(path, self)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{VectorElementType, Distance};
use itertools::Itertools;
use rand::seq::SliceRandom;
use rand::thread_rng;
use crate::fixtures::index_fixtures::{TestRawScorerProducer, FakeConditionChecker, random_vector};
use std::fs::File;
use std::io::Write;
use ndarray::Array;
use tempdir::TempDir;
#[test]
fn test_connect_new_point() {
let num_points = 10;
let m = 6;
let ef_construct = 32;
// See illustration in docs
let points: Vec<Vec<VectorElementType>> = vec![
vec![21.79, 7.18], // Target
vec![20.58, 5.46], // 1 B - yes
vec![21.19, 4.51], // 2 C
vec![24.73, 8.24], // 3 D - yes
vec![24.55, 9.98], // 4 E
vec![26.11, 6.85], // 5 F
vec![17.64, 11.14], // 6 G - yes
vec![14.97, 11.52], // 7 I
vec![14.97, 9.60], // 8 J
vec![16.23, 14.32], // 9 H
vec![12.69, 19.13], // 10 K
];
let scorer = |a: PointOffsetType, b: PointOffsetType| {
-(
(points[a as usize][0] - points[b as usize][0]).powi(2) +
(points[a as usize][1] - points[b as usize][1]).powi(2)
).sqrt()
};
let mut insert_ids = (1..points.len() as PointOffsetType).collect_vec();
let mut candidates = FixedLengthPriorityQueue::new(insert_ids.len());
for id in insert_ids.iter().cloned() {
candidates.push(ScoredPointOffset {
idx: id,
score: scorer(0, id),
});
}
let res = GraphLayers::select_candidates_with_heuristic(
candidates, m, scorer,
);
assert_eq!(&res, &vec![1, 3, 6]);
let mut graph_layers = GraphLayers::new(num_points, m, m, ef_construct, 1, true);
insert_ids.shuffle(&mut thread_rng());
for id in insert_ids.iter().cloned() {
let level_m = graph_layers.get_m(0);
GraphLayers::connect_new_point(
&mut graph_layers.links_layers[0][0],
id,
0,
level_m,
scorer,
)
}
assert_eq!(graph_layers.links(0, 0), &vec![1, 2, 3, 4, 5, 6]);
}
fn search_in_graph(query: &Vec<VectorElementType>, top: usize, vector_storage: &TestRawScorerProducer, graph: &GraphLayers) -> Vec<ScoredPointOffset> {
let fake_condition_checker = FakeConditionChecker {};
let raw_scorer = vector_storage.get_raw_scorer(query.clone());
let scorer = FilteredScorer {
raw_scorer: &raw_scorer,
condition_checker: &fake_condition_checker,
filter: None,
};
let ef = 16;
graph.search(top, ef, &scorer)
}
const M: usize = 8;
fn create_graph_layer(num_vectors: usize, dim: usize, use_heuristic: bool) -> (TestRawScorerProducer, GraphLayers) {
let m = M;
let ef_construct = 16;
let entry_points_num = 10;
let vector_holder = TestRawScorerProducer::new(dim, num_vectors, Distance::Cosine);
let mut graph_layers = GraphLayers::new(
num_vectors, m, m * 2, ef_construct, entry_points_num, use_heuristic,
);
let mut rng = thread_rng();
for idx in 0..(num_vectors as PointOffsetType) {
let fake_condition_checker = FakeConditionChecker {};
let added_vector = vector_holder.vectors[idx as usize].to_vec();
let raw_scorer = vector_holder.get_raw_scorer(added_vector.clone());
let scorer = FilteredScorer {
raw_scorer: &raw_scorer,
condition_checker: &fake_condition_checker,
filter: None,
};
let level = graph_layers.get_random_layer(&mut rng);
graph_layers.link_new_point(idx, level, &scorer);
}
(vector_holder, graph_layers)
}
#[test]
fn test_search_on_level() {
let dim = 8;
let m = 8;
let ef_construct = 32;
let entry_points_num = 10;
let num_vectors = 10;
let vector_holder = TestRawScorerProducer::new(dim, num_vectors, Distance::Dot);
let mut graph_layers = GraphLayers::new(
num_vectors, m, m * 2, ef_construct, entry_points_num, false,
);
graph_layers.links_layers[0][0] = vec![1, 2, 3, 4, 5, 6];
let linking_idx: PointOffsetType = 7;
let fake_condition_checker = FakeConditionChecker {};
let added_vector = vector_holder.vectors[linking_idx as usize].to_vec();
let raw_scorer = vector_holder.get_raw_scorer(added_vector);
let scorer = FilteredScorer {
raw_scorer: &raw_scorer,
condition_checker: &fake_condition_checker,
filter: None,
};
let nearest_on_level = graph_layers.search_on_level(
ScoredPointOffset {
idx: 0,
score: scorer.score_point(0),
},
0,
32,
&scorer,
);
assert_eq!(nearest_on_level.len(), graph_layers.links_layers[0][0].len() + 1);
for nearest in nearest_on_level.iter() {
// eprintln!("nearest = {:#?}", nearest);
assert_eq!(nearest.score, scorer.score_internal(linking_idx, nearest.idx))
}
}
#[test]
fn test_save_and_load() {
let num_vectors = 100;
let dim = 8;
let top = 5;
let (vector_holder, graph_layers) = create_graph_layer(num_vectors, dim, false);
let mut rng = thread_rng();
let query = random_vector(&mut rng, dim);
let res1 = search_in_graph(&query, top, &vector_holder, &graph_layers);
let dir = TempDir::new("graph_dir").unwrap();
let path = GraphLayers::get_path(dir.path());
graph_layers.save(&path).unwrap();
let graph2 = GraphLayers::load(&path).unwrap();
let res2 = search_in_graph(&query, top, &vector_holder, &graph2);
assert_eq!(res1, res2)
}
#[test]
fn test_add_points() {
let num_vectors = 1000;
let dim = 8;
let (vector_holder, graph_layers) = create_graph_layer(num_vectors, dim, false);
let mut rng = thread_rng();
let main_entry = graph_layers.entry_points.get_entry_point(|_x| true)
.expect("Expect entry point to exists");
assert!(main_entry.level > 0);
let num_levels = graph_layers.links_layers
.iter()
.map(|x| x.len())
.max().unwrap();
assert_eq!(main_entry.level + 1, num_levels);
let total_links_0: usize = graph_layers.links_layers
.iter()
.map(|x| x[0].len()).sum();
assert!(total_links_0 > 0);
assert!(total_links_0 as f64 / num_vectors as f64 > M as f64);
let top = 5;
let query = random_vector(&mut rng, dim);
let processed_query = Array::from(vector_holder.metric.preprocess(query.clone()));
let mut reference_top = FixedLengthPriorityQueue::new(top);
for (idx, vec) in vector_holder.vectors.iter().enumerate() {
reference_top.push(
ScoredPointOffset {
idx: idx as PointOffsetType,
score: vector_holder.metric.blas_similarity(vec, &processed_query),
}
);
}
let graph_search = search_in_graph(&query, top, &vector_holder, &graph_layers);
assert_eq!(reference_top.into_vec(), graph_search);
}
#[test]
#[ignore]
fn test_draw_hnsw_graph() {
let dim = 2;
let num_vectors = 500;
let (vector_holder, graph_layers) = create_graph_layer(num_vectors, dim, true);
let graph_json = serde_json::to_string_pretty(&graph_layers).unwrap();
let vectors_json = serde_json::to_string_pretty(&vector_holder.vectors.iter().map(|x| x.to_vec()).collect_vec()).unwrap();
let mut file = File::create("graph.json").unwrap();
file.write_all(format!("{{ \"graph\": {}, \n \"vectors\": {} }}", graph_json, vectors_json).as_bytes()).unwrap();
}
}

244
lib/segment/src/index/hnsw_index/hnsw.rs Normal file
View File

@@ -0,0 +1,244 @@
use crate::entry::entry_point::OperationResult;
use std::path::{Path, PathBuf};
use std::fs::create_dir_all;
use crate::index::index::{VectorIndex, PayloadIndex};
use crate::types::{SearchParams, Filter, PointOffsetType, VectorElementType, FieldCondition, HnswConfig};
use crate::vector_storage::vector_storage::{ScoredPointOffset, VectorStorage};
use std::sync::Arc;
use atomic_refcell::AtomicRefCell;
use crate::payload_storage::payload_storage::ConditionChecker;
use std::cmp::max;
use std::ops::Deref;
use crate::index::hnsw_index::point_scorer::FilteredScorer;
use rand::thread_rng;
use rand::prelude::ThreadRng;
use crate::index::hnsw_index::config::HnswGraphConfig;
use crate::index::hnsw_index::graph_layers::GraphLayers;
use crate::types::Condition::Field;
use crate::index::hnsw_index::build_condition_checker::BuildConditionChecker;
use crate::index::sample_estimation::sample_check_cardinality;
const HNSW_USE_HEURISTIC: bool = true;
pub struct HNSWIndex {
condition_checker: Arc<AtomicRefCell<dyn ConditionChecker>>,
vector_storage: Arc<AtomicRefCell<dyn VectorStorage>>,
payload_index: Arc<AtomicRefCell<dyn PayloadIndex>>,
config: HnswGraphConfig,
path: PathBuf,
thread_rng: ThreadRng,
graph: GraphLayers,
}
impl HNSWIndex {
pub fn open(
path: &Path,
condition_checker: Arc<AtomicRefCell<dyn ConditionChecker>>,
vector_storage: Arc<AtomicRefCell<dyn VectorStorage>>,
payload_index: Arc<AtomicRefCell<dyn PayloadIndex>>,
hnsw_config: HnswConfig,
) -> OperationResult<Self> {
create_dir_all(path)?;
let rng = thread_rng();
let config_path = HnswGraphConfig::get_config_path(path);
let config = if config_path.exists() {
HnswGraphConfig::load(&config_path)?
} else {
HnswGraphConfig::new(hnsw_config.m, hnsw_config.ef_construct, hnsw_config.full_scan_threshold)
};
let graph_path = GraphLayers::get_path(path);
let graph = if graph_path.exists() {
GraphLayers::load(graph_path.as_path())?
} else {
let total_points = vector_storage.borrow().total_vector_count();
GraphLayers::new(
vector_storage.borrow().total_vector_count(),
config.m,
config.m0,
config.ef_construct,
max(1, total_points / hnsw_config.full_scan_threshold * 10),
HNSW_USE_HEURISTIC,
)
};
Ok(HNSWIndex {
condition_checker,
vector_storage,
payload_index,
config,
path: path.to_owned(),
thread_rng: rng,
graph,
})
}
fn save_config(&self) -> OperationResult<()> {
let config_path = HnswGraphConfig::get_config_path(self.path.as_path());
self.config.save(&config_path)
}
fn save_graph(&self) -> OperationResult<()> {
let graph_path = GraphLayers::get_path(self.path.as_path());
self.graph.save(&graph_path)
}
pub fn save(&self) -> OperationResult<()> {
self.save_config()?;
self.save_graph()?;
Ok(())
}
pub fn link_point(&mut self, point_id: PointOffsetType, points_scorer: &FilteredScorer) {
let point_level = self.graph.get_random_layer(&mut self.thread_rng);
self.graph.link_new_point(point_id, point_level, points_scorer);
}
pub fn build_filtered_graph(&self, condition: FieldCondition, block_condition_checker: &mut BuildConditionChecker) -> GraphLayers {
block_condition_checker.filter_list.next_iteration();
let filter = Filter::new_must(Field(condition));
let payload_index = self.payload_index.borrow();
let vector_storage = self.vector_storage.borrow();
for block_point_id in payload_index.query_points(&filter) {
block_condition_checker.filter_list.check_and_update_visited(block_point_id);
}
let mut graph = GraphLayers::new(
self.vector_storage.borrow().total_vector_count(),
self.config.m,
self.config.m0,
self.config.ef_construct,
1,
HNSW_USE_HEURISTIC,
);
for block_point_id in payload_index.query_points(&filter) {
let vector = vector_storage.get_vector(block_point_id).unwrap();
let raw_scorer = vector_storage.raw_scorer(vector);
let points_scorer = FilteredScorer {
raw_scorer: raw_scorer.as_ref(),
condition_checker: block_condition_checker,
filter: None,
};
let level = self.graph.point_level(block_point_id);
graph.link_new_point(block_point_id, level, &points_scorer);
}
graph
}
pub fn search_with_graph(&self, vector: &Vec<VectorElementType>, filter: Option<&Filter>, top: usize, params: Option<&SearchParams>) -> Vec<ScoredPointOffset> {
let req_ef = params.and_then(|params| params.hnsw_ef).unwrap_or(self.config.ef);
// ef should always be bigger that required top
let ef = max(req_ef, top);
let vector_storage = self.vector_storage.borrow();
let raw_scorer = vector_storage.raw_scorer(vector.clone());
let condition_checker = self.condition_checker.borrow();
let points_scorer = FilteredScorer {
raw_scorer: raw_scorer.as_ref(),
condition_checker: condition_checker.deref(),
filter,
};
self.graph.search(top, ef, &points_scorer)
}
}
impl VectorIndex for HNSWIndex {
fn search(&self, vector: &Vec<VectorElementType>, filter: Option<&Filter>, top: usize, params: Option<&SearchParams>) -> Vec<ScoredPointOffset> {
match filter {
None => self.search_with_graph(vector, None, top, params),
Some(query_filter) => {
// depending on the amount of filtered-out points the optimal strategy could be
// - to retrieve possible points and score them after
// - to use HNSW index with filtering condition
let payload_index = self.payload_index.borrow();
let query_cardinality = payload_index.estimate_cardinality(query_filter);
let vector_storage = self.vector_storage.borrow();
if query_cardinality.max < self.config.indexing_threshold {
// if cardinality is small - use plain index
let mut filtered_ids = payload_index.query_points(query_filter);
return vector_storage.score_points(vector, &mut filtered_ids, top);
}
if query_cardinality.min > self.config.indexing_threshold {
// if cardinality is high enough - use HNSW index
return self.search_with_graph(vector, filter, top, params);
}
// Fast cardinality estimation is not enough, do sample estimation of cardinality
let condition_checker = self.condition_checker.borrow();
return if sample_check_cardinality(
vector_storage.sample_ids(),
|idx| condition_checker.check(idx, query_filter),
self.config.indexing_threshold,
vector_storage.vector_count()
) {
// if cardinality is high enough - use HNSW index
self.search_with_graph(vector, filter, top, params)
} else {
// if cardinality is small - use plain index
let mut filtered_ids = payload_index.query_points(query_filter);
vector_storage.score_points(vector, &mut filtered_ids, top)
};
}
}
}
fn build_index(&mut self) -> OperationResult<()> {
// Build main index graph
let vector_storage = self.vector_storage.borrow();
let condition_checker = self.condition_checker.borrow();
let mut rng = thread_rng();
let total_points = vector_storage.total_vector_count();
self.graph = GraphLayers::new(
total_points,
self.config.m,
self.config.m0,
self.config.ef_construct,
max(1, total_points / self.config.indexing_threshold * 10),
HNSW_USE_HEURISTIC,
);
for vector_id in vector_storage.iter_ids() {
let vector = vector_storage.get_vector(vector_id).unwrap();
let raw_scorer = vector_storage.raw_scorer(vector);
let points_scorer = FilteredScorer {
raw_scorer: raw_scorer.as_ref(),
condition_checker: condition_checker.deref(),
filter: None,
};
let level = self.graph.get_random_layer(&mut rng);
self.graph.link_new_point(vector_id, level, &points_scorer);
}
let total_vectors_count = vector_storage.total_vector_count();
let mut block_condition_checker = BuildConditionChecker::new(total_vectors_count);
let payload_index = self.payload_index.borrow();
// ToDo: Think about using connectivity threshold (based on m0) instead of `indexing_threshold`
for payload_block in payload_index.payload_blocks(self.config.indexing_threshold) {
let block_graph = self.build_filtered_graph(payload_block.condition, &mut block_condition_checker);
self.graph.merge_from_other(block_graph);
}
self.save()
}
}

11
lib/segment/src/index/hnsw_index/mod.rs Normal file
View File

@@ -0,0 +1,11 @@
pub mod hnsw;
pub mod graph_layers;
pub mod point_scorer;
mod config;
mod entry_points;
mod search_context;
mod build_cache;
pub mod build_condition_checker;
#[cfg(test)]
mod tests;

48
lib/segment/src/index/hnsw_index/point_scorer.rs Normal file
View File

@@ -0,0 +1,48 @@
use crate::types::{PointOffsetType, Filter, ScoreType};
use crate::vector_storage::vector_storage::{ScoredPointOffset, RawScorer};
use crate::payload_storage::payload_storage::ConditionChecker;
pub struct FilteredScorer<'a> {
pub raw_scorer: &'a dyn RawScorer,
pub condition_checker: &'a dyn ConditionChecker,
pub filter: Option<&'a Filter>,
}
impl FilteredScorer<'_> {
pub fn check_point(&self, point_id: PointOffsetType) -> bool {
match self.filter {
None => self.raw_scorer.check_point(point_id),
Some(f) => self.condition_checker.check(point_id, f) && self.raw_scorer.check_point(point_id)
}
}
pub fn score_iterable_points<F>(&self, points_iterator: &mut dyn Iterator<Item=PointOffsetType>, limit: usize, action: F)
where F: FnMut(ScoredPointOffset) {
match self.filter {
None => self.raw_scorer.score_points(points_iterator).take(limit).for_each(action),
Some(f) => {
let mut points_filtered_iterator = points_iterator
.filter(move |id| self.condition_checker.check(*id, f));
self.raw_scorer.score_points(&mut points_filtered_iterator).take(limit).for_each(action)
}
};
}
pub fn score_points<F>(&self, ids: &[PointOffsetType], limit: usize, action: F)
where F: FnMut(ScoredPointOffset) {
let mut points_iterator = ids
.iter()
.cloned();
self.score_iterable_points(&mut points_iterator, limit, action);
}
pub fn score_point(&self, point_id: PointOffsetType) -> ScoreType {
self.raw_scorer.score_point(point_id)
}
pub fn score_internal(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> ScoreType {
self.raw_scorer.score_internal(point_a, point_b)
}
}

35
lib/segment/src/index/hnsw_index/search_context.rs Normal file
View File

@@ -0,0 +1,35 @@
use crate::spaces::tools::FixedLengthPriorityQueue;
use crate::vector_storage::vector_storage::ScoredPointOffset;
use crate::types::PointOffsetType;
/// Structure that holds context of the search
pub struct SearchContext {
pub nearest: FixedLengthPriorityQueue<ScoredPointOffset>,
pub candidates: Vec<PointOffsetType>,
}
impl SearchContext {
pub fn new(entry_point: ScoredPointOffset, ef: usize) -> Self {
let mut nearest = FixedLengthPriorityQueue::new(ef);
nearest.push(entry_point);
SearchContext {
nearest,
candidates: vec![entry_point.idx]
}
}
/// Updates search context with new scored point.
/// If it is closer than existing - also add it to candidates for further search
pub fn process_candidate(&mut self, score_point: ScoredPointOffset) {
let was_added = match self.nearest.push(score_point.clone()) {
None => true,
Some(removed) => removed.idx != score_point.idx
};
if was_added {
self.candidates.push(score_point.idx)
// ToDo: update cache here
}
}
}

0
lib/segment/src/index/hnsw_index/tests/mod.rs Normal file
View File

10
lib/segment/src/index/index.rs
View File

@@ -1,10 +1,10 @@
use crate::types::{Filter, PointOffsetType, VectorElementType, SearchParams, PayloadKeyType};
use crate::vector_storage::vector_storage::ScoredPointOffset;
use crate::entry::entry_point::OperationResult;
use crate::index::field_index::CardinalityEstimation;
use crate::index::field_index::{CardinalityEstimation, PayloadBlockCondition};
/// Trait for vector searching
pub trait Index {
pub trait VectorIndex {
/// Return list of Ids with fitting
fn search(&self,
vector: &Vec<VectorElementType>,
@@ -32,5 +32,9 @@ pub trait PayloadIndex {
fn estimate_cardinality(&self, query: &Filter) -> CardinalityEstimation;
/// Return list of all point ids, which satisfy filtering criteria
fn query_points(&self, query: &Filter) -> Box<dyn Iterator<Item=PointOffsetType> + '_>;
fn query_points<'a>(&'a self, query: &'a Filter) -> Box<dyn Iterator<Item=PointOffsetType> + 'a>;
/// Iterate conditions for payload blocks with minimum size of `threshold`
/// Required for building HNSW index
fn payload_blocks(&self, threshold: usize) -> Box<dyn Iterator<Item=PayloadBlockCondition> + '_>;
}

7
lib/segment/src/index/mod.rs
View File

@@ -1,6 +1,11 @@
pub mod plain_payload_index;
pub mod index;
pub mod struct_payload_index;
pub mod query_estimator;
pub mod hnsw_index;
mod field_index;
mod payload_config;
pub mod query_estimator;
mod visited_pool;
mod sample_estimation;

46
lib/segment/src/index/plain_payload_index.rs
View File

@@ -1,6 +1,6 @@
use crate::vector_storage::vector_storage::{ScoredPointOffset, VectorStorage};
use crate::index::index::{Index, PayloadIndex};
use crate::types::{Filter, VectorElementType, Distance, SearchParams, PointOffsetType, PayloadKeyType};
use crate::index::index::{VectorIndex, PayloadIndex};
use crate::types::{Filter, VectorElementType, SearchParams, PointOffsetType, PayloadKeyType};
use crate::payload_storage::payload_storage::{ConditionChecker};
use std::sync::Arc;
@@ -9,8 +9,7 @@ use crate::entry::entry_point::OperationResult;
use crate::index::payload_config::PayloadConfig;
use std::path::{Path, PathBuf};
use std::fs::create_dir_all;
use crate::index::field_index::CardinalityEstimation;
use itertools::Itertools;
use crate::index::field_index::{CardinalityEstimation, PayloadBlockCondition};
pub struct PlainPayloadIndex {
@@ -79,23 +78,17 @@ impl PayloadIndex for PlainPayloadIndex {
self.save_config()
}
fn estimate_cardinality(&self, query: &Filter) -> CardinalityEstimation {
let mut matched_points = 0;
let condition_checker = self.condition_checker.borrow();
for i in self.vector_storage.borrow().iter_ids() {
if condition_checker.check(i, query) {
matched_points += 1;
}
}
fn estimate_cardinality(&self, _query: &Filter) -> CardinalityEstimation {
let total_points = self.vector_storage.borrow().vector_count();
CardinalityEstimation {
primary_clauses: vec![],
min: matched_points,
exp: matched_points,
max: matched_points
min: 0,
exp: total_points / 2,
max: total_points
}
}
fn query_points(&self, query: &Filter) -> Box<dyn Iterator<Item=PointOffsetType> + '_> {
fn query_points<'a>(&'a self, query: &'a Filter) -> Box<dyn Iterator<Item=PointOffsetType> + 'a> {
let mut matched_points = vec![];
let condition_checker = self.condition_checker.borrow();
for i in self.vector_storage.borrow().iter_ids() {
@@ -105,31 +98,33 @@ impl PayloadIndex for PlainPayloadIndex {
}
return Box::new(matched_points.into_iter());
}
fn payload_blocks(&self, _threshold: usize) -> Box<dyn Iterator<Item=PayloadBlockCondition> + '_> {
// No blocks for un-indexed payload
Box::new(vec![].into_iter())
}
}
pub struct PlainIndex {
vector_storage: Arc<AtomicRefCell<dyn VectorStorage>>,
payload_index: Arc<AtomicRefCell<dyn PayloadIndex>>,
distance: Distance,
}
impl PlainIndex {
pub fn new(
vector_storage: Arc<AtomicRefCell<dyn VectorStorage>>,
payload_index: Arc<AtomicRefCell<dyn PayloadIndex>>,
distance: Distance,
payload_index: Arc<AtomicRefCell<dyn PayloadIndex>>
) -> PlainIndex {
return PlainIndex {
vector_storage,
payload_index,
distance,
payload_index
};
}
}
impl Index for PlainIndex {
impl VectorIndex for PlainIndex {
fn search(
&self,
vector: &Vec<VectorElementType>,
@@ -139,10 +134,11 @@ impl Index for PlainIndex {
) -> Vec<ScoredPointOffset> {
match filter {
Some(filter) => {
let filtered_ids = self.payload_index.borrow().query_points(filter).collect_vec();
self.vector_storage.borrow().score_points(vector, &filtered_ids, top, &self.distance)
let borrowed_payload_index = self.payload_index.borrow();
let mut filtered_ids = borrowed_payload_index.query_points(filter);
self.vector_storage.borrow().score_points(vector, &mut filtered_ids, top)
}
None => self.vector_storage.borrow().score_all(vector, top, &self.distance)
None => self.vector_storage.borrow().score_all(vector, top)
}
}

4
lib/segment/src/index/query_estimator.rs
View File

@@ -126,7 +126,7 @@ fn estimate_must_not<F>(estimator: &F, conditions: &Vec<Condition>, total: usize
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{FieldCondition, HasIdCondition};
use crate::types::{FieldCondition, HasIdCondition, PointOffsetType};
const TOTAL: usize = 1000;
@@ -165,7 +165,7 @@ mod tests {
_ => CardinalityEstimation::unknown(TOTAL)
},
Condition::HasId(has_id) => CardinalityEstimation {
primary_clauses: vec![PrimaryCondition::Ids(has_id.has_id.iter().map(|x| *x as usize).collect())],
primary_clauses: vec![PrimaryCondition::Ids(has_id.has_id.iter().map(|x| *x as PointOffsetType).collect())],
min: has_id.has_id.len(),
exp: has_id.has_id.len(),
max: has_id.has_id.len(),

102
lib/segment/src/index/sample_estimation.rs Normal file
View File

@@ -0,0 +1,102 @@
use std::cmp::{min, max};
use crate::types::{PointOffsetType};
const MAX_ESTIMATED_POINTS: usize = 1000;
/// How many points do we need to check in order to estimate expected query cardinality.
/// Based on https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval
#[allow(dead_code)]
fn estimate_required_sample_size(total: usize, confidence_interval: usize) -> usize {
let confidence_interval = min(confidence_interval, total);
let z = 1.96; // percentile 0.95 of normal distribution
let index_fraction = confidence_interval as f64 / total as f64 / 2.0;
let h = 0.5; // success rate which requires most number of estimations
let estimated_size = h * (1. - h) / (index_fraction / z).powi(2);
return max(estimated_size as usize, 10);
}
/// Returns (expected cardinality ± confidence interval at 0.99)
/// Based on https://en.wikipedia.org/wiki/Binomial_proportion_confidence_interval#Agresti%E2%80%93Coull_interval
fn confidence_agresti_coull_interval(trials: usize, positive: usize, total: usize) -> (i64, i64) {
let z = 2.; // heuristics
let nhat = trials as f64 + z * z;
let phat = (positive as f64 + z * z / 2.) / nhat;
let interval = z * ((phat / nhat) * (1. - phat)).sqrt();
let expected = (phat * total as f64) as i64;
let delta = (interval * total as f64) as i64;
return (expected, delta);
}
/// Tests if given `query` have cardinality higher than the `threshold`
/// Iteratively samples points until the decision could be made with confidence
pub fn sample_check_cardinality(
sample_points: impl Iterator<Item=PointOffsetType>,
checker: impl Fn(PointOffsetType) -> bool,
threshold: usize,
total_points: usize
) -> bool {
let mut matched_points = 0;
let mut total_checked = 0;
let mut exp = 0;
let mut interval ;
for idx in sample_points.take(MAX_ESTIMATED_POINTS) {
matched_points += checker(idx) as usize;
total_checked += 1;
let estimation = confidence_agresti_coull_interval(total_checked, matched_points, total_points);
exp = estimation.0;
interval = estimation.1;
if exp - interval > threshold as i64 {
return true
}
if exp + interval < threshold as i64 {
return false
}
}
exp > threshold as i64
}
#[cfg(test)]
mod tests {
use super::*;
use rand::Rng;
#[test]
fn test_confidence_interval() {
let total = 100_000;
let true_p = 0.25;
let mut rng = rand::thread_rng();
let mut delta = 100_000;
let mut positive = 0;
for i in 1..=101 {
positive += rng.gen_bool(true_p) as usize;
if i % 20 == 1 {
let interval = confidence_agresti_coull_interval(i, positive, total);
assert!(interval.1 < delta);
delta = interval.1;
eprintln!("confidence_agresti_coull_interval({}, {}, {}) = {:#?}", i, positive, total, interval);
}
}
}
#[test]
fn test_sample_check_cardinality() {
let res = sample_check_cardinality(
vec![1,2,3,4,5,6,7,8,9,10,11,12].into_iter(),
|idx| idx % 2 == 0,
10_000,
100_000
);
assert!(res)
}
}

54
lib/segment/src/index/struct_payload_index.rs
View File

@@ -14,10 +14,11 @@ use crate::index::index::PayloadIndex;
use crate::index::payload_config::PayloadConfig;
use crate::payload_storage::payload_storage::{ConditionChecker, PayloadStorage};
use crate::types::{Filter, PayloadKeyType, FieldCondition, Condition, PointOffsetType};
use crate::index::field_index::{CardinalityEstimation, PrimaryCondition};
use crate::index::query_estimator::estimate_filter;
use crate::index::field_index::{CardinalityEstimation, PrimaryCondition, PayloadBlockCondition};
use crate::index::query_estimator::{estimate_filter};
use crate::vector_storage::vector_storage::VectorStorage;
use crate::id_mapper::id_mapper::IdMapper;
use crate::index::visited_pool::VisitedPool;
pub const PAYLOAD_FIELD_INDEX_PATH: &str = "fields";
@@ -31,6 +32,7 @@ pub struct StructPayloadIndex {
field_indexes: IndexesMap,
config: PayloadConfig,
path: PathBuf,
visited_pool: VisitedPool,
}
impl StructPayloadIndex {
@@ -143,7 +145,8 @@ impl StructPayloadIndex {
id_mapper,
field_indexes: Default::default(),
config,
path: path.to_owned()
path: path.to_owned(),
visited_pool: Default::default(),
};
if !index.config_path().exists() {
@@ -241,7 +244,7 @@ impl PayloadIndex for StructPayloadIndex {
}
fn estimate_cardinality(&self, query: &Filter) -> CardinalityEstimation {
let total = self.total_points();
let total_points = self.total_points();
let estimator = |condition: &Condition| {
match condition {
@@ -265,16 +268,30 @@ impl PayloadIndex for StructPayloadIndex {
}
};
estimate_filter(&estimator, query, total)
estimate_filter(&estimator, query, total_points)
}
fn query_points(&self, query: &Filter) -> Box<dyn Iterator<Item=PointOffsetType> + '_> {
fn payload_blocks(&self, threshold: usize) -> Box<dyn Iterator<Item=PayloadBlockCondition> + '_> {
let iter = self.field_indexes
.iter()
.map(move |(key, indexes)| {
indexes
.iter()
.map(move |field_index| field_index.payload_blocks(threshold, key.clone()))
.flatten()
}).flatten();
Box::new(iter)
}
fn query_points<'a>(&'a self, query: &'a Filter) -> Box<dyn Iterator<Item=PointOffsetType> + 'a> {
// Assume query is already estimated to be small enough so we can iterate over all matched ids
let query_cardinality = self.estimate_cardinality(query);
let condition_checker = self.condition_checker.borrow();
let vector_storage_ref = self.vector_storage.borrow();
let full_scan_iterator = vector_storage_ref.iter_ids(); // Should not be used if filter restricted by indexed fields
let condition_checker = self.condition_checker.borrow();
let query_cardinality = self.estimate_cardinality(query);
return if query_cardinality.primary_clauses.is_empty() {
let full_scan_iterator = vector_storage_ref.iter_ids();
// Worst case: query expected to return few matches, but index can't be used
let matched_points = full_scan_iterator
.filter(|i| condition_checker.check(*i, query))
@@ -283,7 +300,10 @@ impl PayloadIndex for StructPayloadIndex {
Box::new(matched_points.into_iter())
} else {
// CPU-optimized strategy here: points are made unique before applying other filters.
let preselected: HashSet<PointOffsetType> = query_cardinality.primary_clauses.iter()
// ToDo: Implement iterator which holds the `visited_pool` and borrowed `vector_storage_ref` to prevent `preselected` array creation
let mut visited_list = self.visited_pool.get(vector_storage_ref.total_vector_count());
let preselected: Vec<PointOffsetType> = query_cardinality.primary_clauses.iter()
.map(|clause| {
match clause {
PrimaryCondition::Condition(field_condition) => self.query_field(field_condition)
@@ -291,12 +311,16 @@ impl PayloadIndex for StructPayloadIndex {
PrimaryCondition::Ids(ids) => Box::new(ids.iter().cloned())
}
})
.flat_map(|x| x)
.flatten()
.filter(|id| !visited_list.check_and_update_visited(*id))
.filter(move |i| condition_checker.check(*i, query))
.collect();
let matched_points = preselected.into_iter()
.filter(|i| condition_checker.check(*i, query))
.collect_vec();
Box::new(matched_points.into_iter())
self.visited_pool.return_back(visited_list);
let matched_points_iter = preselected.into_iter();
Box::new(matched_points_iter)
};
}
}

86
lib/segment/src/index/visited_pool.rs Normal file
View File

@@ -0,0 +1,86 @@
use crate::types::PointOffsetType;
use parking_lot::RwLock;
/// Max number of visited lists to preserve in memory
/// If more than this number of concurrent requests occurred - new list will be created dynamically,
/// but will be deleted right after query finishes.
/// Implemented in order to limit memory leak
const POOL_KEEP_LIMIT: usize = 16;
#[derive(Debug)]
pub struct VisitedList {
current_iter: usize,
visit_counters: Vec<usize>,
}
impl VisitedList {
pub fn new(num_points: usize) -> Self {
VisitedList {
current_iter: 1,
visit_counters: vec![0; num_points],
}
}
/// Return `true` if visited
pub fn check(&self, point_id: PointOffsetType) -> bool {
self.visit_counters
.get(point_id as usize)
.map(|x| *x >= self.current_iter)
.unwrap_or(false)
}
/// Updates visited list
/// return `true` if point was visited before
pub fn check_and_update_visited(&mut self, point_id: PointOffsetType) -> bool {
let idx = point_id as usize;
if idx >= self.visit_counters.len() {
self.visit_counters.resize(idx + 1, 0);
}
let prev_value = self.visit_counters[idx];
self.visit_counters[idx] = self.current_iter;
prev_value >= self.current_iter
}
pub fn next_iteration(&mut self) {
self.current_iter += 1;
}
}
#[derive(Debug)]
pub struct VisitedPool {
pool: RwLock<Vec<VisitedList>>
}
impl VisitedPool {
pub fn new() -> Self {
VisitedPool {
pool: RwLock::new(vec![])
}
}
pub fn get(&self, num_points: usize) -> VisitedList {
match self.pool.write().pop() {
None => VisitedList::new(num_points),
Some(mut vl) => {
vl.next_iteration();
vl
}
}
}
pub fn return_back(&self, visited_list: VisitedList) {
let mut pool = self.pool.write();
if pool.len() < POOL_KEEP_LIMIT {
pool.push(visited_list)
}
}
}
impl Default for VisitedPool {
fn default() -> Self {
VisitedPool::new()
}
}

6
lib/segment/src/lib.rs
View File

@@ -1,13 +1,13 @@
mod id_mapper;
mod query_planner;
mod index;
mod payload_storage;
pub mod payload_storage;
pub mod index;
pub mod vector_storage;
pub mod segment;
pub mod spaces;
pub mod segment_constructor;
pub mod entry;
pub mod types;
pub mod fixtures;
mod common;

2
lib/segment/src/payload_storage/mod.rs
View File

@@ -1,6 +1,6 @@
pub mod query_checker;
pub mod simple_payload_storage;
pub mod payload_storage;
mod condition_checker;
pub mod condition_checker;

2
lib/segment/src/payload_storage/payload_storage.rs
View File

@@ -75,6 +75,6 @@ pub trait PayloadStorage {
pub trait ConditionChecker {
/// Check if point satisfies filter condition
/// Check if point satisfies filter condition. Return true if satisfies
fn check(&self, point_id: PointOffsetType, query: &Filter) -> bool;
}

2
lib/segment/src/query_planner/mod.rs
View File

@@ -1,2 +0,0 @@
pub mod simple_query_planner;
pub mod query_planner;

18
lib/segment/src/query_planner/query_planner.rs
View File

@@ -1,18 +0,0 @@
use crate::types::{VectorElementType, Filter, SearchParams};
use crate::vector_storage::vector_storage::ScoredPointOffset;
use crate::entry::entry_point::OperationResult;
/// Similar to `Index`, but should operate with multiple possible indexes + post-filtering
pub trait QueryPlanner {
/// Performs search of vector in the most efficient way according to heuristics
fn search(&self,
vector: &Vec<VectorElementType>,
filter: Option<&Filter>,
top: usize,
params: Option<&SearchParams>,
) -> Vec<ScoredPointOffset>;
/// Force internal index rebuild.
fn build_index(&mut self) -> OperationResult<()>;
}

35
lib/segment/src/query_planner/simple_query_planner.rs
View File

@@ -1,35 +0,0 @@
use crate::index::index::Index;
use crate::query_planner::query_planner::QueryPlanner;
use crate::types::{Filter, VectorElementType, SearchParams};
use crate::vector_storage::vector_storage::ScoredPointOffset;
use atomic_refcell::AtomicRefCell;
use std::sync::Arc;
use crate::entry::entry_point::OperationResult;
pub struct SimpleQueryPlanner {
index: Arc<AtomicRefCell<dyn Index>>
}
impl QueryPlanner for SimpleQueryPlanner {
fn search(&self,
vector: &Vec<VectorElementType>,
filter: Option<&Filter>,
top: usize,
params: Option<&SearchParams>,
) -> Vec<ScoredPointOffset> {
self.index.borrow().search(vector, filter, top, params)
}
fn build_index(&mut self) -> OperationResult<()> {
self.index.borrow_mut().build_index()
}
}
impl SimpleQueryPlanner {
pub fn new(index: Arc<AtomicRefCell<dyn Index>>) -> Self {
SimpleQueryPlanner {
index
}
}
}

21
lib/segment/src/segment.rs
View File

@@ -1,16 +1,16 @@
use crate::id_mapper::id_mapper::IdMapper;
use crate::vector_storage::vector_storage::VectorStorage;
use crate::payload_storage::payload_storage::{PayloadStorage};
use crate::payload_storage::payload_storage::{PayloadStorage, ConditionChecker};
use crate::entry::entry_point::{SegmentEntry, OperationResult, OperationError};
use crate::types::{Filter, PayloadKeyType, PayloadType, SeqNumberType, VectorElementType, PointIdType, PointOffsetType, SearchParams, ScoredPoint, TheMap, SegmentInfo, SegmentType, SegmentConfig, SegmentState, PayloadSchemaInfo};
use crate::query_planner::query_planner::QueryPlanner;
use std::sync::{Arc, Mutex};
use atomic_refcell::{AtomicRefCell};
use std::path::PathBuf;
use std::fs::{remove_dir_all};
use std::io::Write;
use atomicwrites::{AtomicFile, AllowOverwrite};
use crate::index::index::PayloadIndex;
use crate::index::index::{PayloadIndex, VectorIndex};
use crate::spaces::tools::mertic_object;
pub const SEGMENT_STATE_FILE: &str = "segment.json";
@@ -24,8 +24,8 @@ pub struct Segment {
pub vector_storage: Arc<AtomicRefCell<dyn VectorStorage>>,
pub payload_storage: Arc<AtomicRefCell<dyn PayloadStorage>>,
pub payload_index: Arc<AtomicRefCell<dyn PayloadIndex>>,
/// User for writing only here.
pub query_planner: Arc<AtomicRefCell<dyn QueryPlanner>>,
pub condition_checker: Arc<AtomicRefCell<dyn ConditionChecker>>,
pub vector_index: Arc<AtomicRefCell<dyn VectorIndex>>,
pub appendable_flag: bool,
pub segment_type: SegmentType,
pub segment_config: SegmentConfig,
@@ -36,7 +36,7 @@ impl Segment {
fn update_vector(&mut self,
old_internal_id: PointOffsetType,
vector: &Vec<VectorElementType>,
vector: Vec<VectorElementType>,
) -> OperationResult<PointOffsetType> {
let new_internal_index = {
let mut vector_storage = self.vector_storage.borrow_mut();
@@ -112,7 +112,7 @@ impl SegmentEntry for Segment {
});
}
let internal_result = self.query_planner.borrow().search(vector, filter, top, params);
let internal_result = self.vector_index.borrow().search(vector, filter, top, params);
let id_mapper = self.id_mapper.borrow();
@@ -139,6 +139,9 @@ impl SegmentEntry for Segment {
return Err(OperationError::WrongVector { expected_dim: vector_dim, received_dim: vector.len() });
}
let metric = mertic_object(&self.segment_config.distance);
let processed_vector = metric.preprocess(vector.clone());
let stored_internal_point = {
let id_mapped = self.id_mapper.borrow();
id_mapped.internal_id(point_id)
@@ -146,9 +149,9 @@ impl SegmentEntry for Segment {
let (was_replaced, new_index) = match stored_internal_point {
Some(existing_internal_id) =>
(true, self.update_vector(existing_internal_id, vector)?),
(true, self.update_vector(existing_internal_id, processed_vector)?),
None =>
(false, self.vector_storage.borrow_mut().put_vector(vector)?)
(false, self.vector_storage.borrow_mut().put_vector(processed_vector)?)
};
self.id_mapper.borrow_mut().set_link(point_id, new_index)?;

2
lib/segment/src/segment_constructor/segment_builder.rs
View File

@@ -81,7 +81,7 @@ impl TryInto<Segment> for SegmentBuilder {
segment.create_field_index(segment.version, field)?;
}
segment.query_planner.borrow_mut().build_index()?;
segment.vector_index.borrow_mut().build_index()?;
segment.flush()?;
// Now segment is going to be evicted from RAM

58
lib/segment/src/segment_constructor/segment_constructor.rs
View File

@@ -3,7 +3,6 @@ use crate::id_mapper::simple_id_mapper::SimpleIdMapper;
use crate::vector_storage::simple_vector_storage::SimpleVectorStorage;
use crate::payload_storage::simple_payload_storage::SimplePayloadStorage;
use crate::index::plain_payload_index::{PlainPayloadIndex, PlainIndex};
use crate::query_planner::simple_query_planner::SimpleQueryPlanner;
use crate::types::{SegmentType, SegmentConfig, Indexes, SegmentState, SeqNumberType, StorageType, PayloadIndexType};
use std::sync::{Arc, Mutex};
use atomic_refcell::AtomicRefCell;
@@ -16,7 +15,8 @@ use std::io::Read;
use crate::vector_storage::memmap_vector_storage::MemmapVectorStorage;
use crate::vector_storage::vector_storage::VectorStorage;
use crate::index::struct_payload_index::StructPayloadIndex;
use crate::index::index::PayloadIndex;
use crate::index::index::{PayloadIndex, VectorIndex};
use crate::index::hnsw_index::hnsw::HNSWIndex;
fn sp<T>(t: T) -> Arc<AtomicRefCell<T>> { Arc::new(AtomicRefCell::new(t)) }
@@ -27,13 +27,22 @@ fn create_segment(version: SeqNumberType, segment_path: &Path, config: &SegmentC
let payload_storage_path = segment_path.join("payload_storage");
let payload_index_path = segment_path.join("payload_index");
let vector_storage_path = segment_path.join("vector_storage");
let vector_index_path = segment_path.join("vector_index");
let id_mapper = sp(SimpleIdMapper::open(mapper_path.as_path())?);
let vector_storage: Arc<AtomicRefCell<dyn VectorStorage>> = match config.storage_type {
StorageType::InMemory => sp(SimpleVectorStorage::open(vector_storage_path.as_path(), config.vector_size)?),
StorageType::Mmap => sp(MemmapVectorStorage::open(vector_storage_path.as_path(), config.vector_size)?),
StorageType::InMemory => sp(SimpleVectorStorage::open(
vector_storage_path.as_path(),
config.vector_size,
config.distance,
)?),
StorageType::Mmap => sp(MemmapVectorStorage::open(
vector_storage_path.as_path(),
config.vector_size,
config.distance,
)?),
};
let payload_storage = sp(SimplePayloadStorage::open(payload_storage_path.as_path())?);
@@ -45,29 +54,31 @@ fn create_segment(version: SeqNumberType, segment_path: &Path, config: &SegmentC
));
let payload_index: Arc<AtomicRefCell<dyn PayloadIndex>> = match config.payload_index.unwrap_or_default() {
PayloadIndexType::Plain => sp(PlainPayloadIndex::open(condition_checker, vector_storage.clone(), &payload_index_path)?),
PayloadIndexType::Plain => sp(PlainPayloadIndex::open(
condition_checker.clone(),
vector_storage.clone(),
&payload_index_path)?),
PayloadIndexType::Struct => sp(StructPayloadIndex::open(
condition_checker,
condition_checker.clone(),
vector_storage.clone(),
payload_storage.clone(),
id_mapper.clone(),
&payload_index_path)?),
};
let index = sp(match config.index {
Indexes::Plain { .. } => PlainIndex::new(
let vector_index: Arc<AtomicRefCell<dyn VectorIndex>> = match config.index {
Indexes::Plain { .. } => sp(PlainIndex::new(
vector_storage.clone(),
payload_index.clone(),
config.distance
),
_ => PlainIndex::new(
)),
Indexes::Hnsw(hnsw_config) => sp(HNSWIndex::open(
&vector_index_path,
condition_checker.clone(),
vector_storage.clone(),
payload_index.clone(),
config.distance
)
// ToDo: Add HNSW index init here
// Indexes::Hnsw { .. } => unimplemented!(),
});
hnsw_config
)?)
};
let segment_type = match config.index {
Indexes::Plain { .. } => match config.payload_index.unwrap_or_default() {
@@ -77,20 +88,19 @@ fn create_segment(version: SeqNumberType, segment_path: &Path, config: &SegmentC
Indexes::Hnsw { .. } => SegmentType::Indexed,
};
let appendable = segment_type == SegmentType::Plain {} && config.storage_type == StorageType::InMemory;
let query_planer = SimpleQueryPlanner::new(index);
let appendable_flag = segment_type == SegmentType::Plain {} && config.storage_type == StorageType::InMemory;
return Ok(Segment {
version,
persisted_version: Arc::new(Mutex::new(version)),
current_path: segment_path.to_owned(),
id_mapper: id_mapper.clone(),
id_mapper,
vector_storage,
payload_storage: payload_storage.clone(),
payload_index: payload_index.clone(),
query_planner: sp(query_planer),
appendable_flag: appendable,
payload_storage,
payload_index,
condition_checker,
vector_index,
appendable_flag,
segment_type,
segment_config: config.clone(),
});

2
lib/segment/src/segment_constructor/simple_segment_constructor.rs
View File

@@ -21,7 +21,7 @@ pub fn build_simple_segment(path: &Path, dim: usize, distance: Distance) -> Oper
index: Indexes::Plain {},
payload_index: None,
distance,
storage_type: Default::default()
storage_type: Default::default(),
},
)
}

28
lib/segment/src/spaces/simple.rs
View File

@@ -1,3 +1,5 @@
extern crate blas_src;
use ndarray::Array1;
use crate::types::{Distance, ScoreType, VectorElementType};
@@ -8,13 +10,34 @@ pub struct DotProductMetric {}
pub struct CosineMetric {}
pub struct EuclidMetric {}
impl Metric for EuclidMetric {
fn distance(&self) -> Distance { Distance::Euclid }
fn similarity(&self, v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
let s: ScoreType = v1.iter().cloned().zip(v2.iter().cloned()).map(|(a, b)| (a - b).powi(2)).sum();
return -s.sqrt();
}
fn blas_similarity(&self, v1: &Array1<VectorElementType>, v2: &Array1<VectorElementType>) -> ScoreType {
let s: ScoreType = v1.iter().cloned().zip(v2.iter().cloned()).map(|(a, b)| (a - b).powi(2)).sum();
return -s.sqrt();
}
fn preprocess(&self, vector: Vec<VectorElementType>) -> Vec<VectorElementType> {
return vector;
}
}
impl Metric for DotProductMetric {
fn distance(&self) -> Distance {
Distance::Dot
}
fn similarity(&self, v1: &[VectorElementType], v2: &[VectorElementType]) -> ScoreType {
let ip: f32 = v1.iter().zip(v2).map(|(a, b)| a * b).sum();
let ip: ScoreType = v1.iter().zip(v2).map(|(a, b)| a * b).sum();
return ip;
}
@@ -42,7 +65,8 @@ impl Metric for CosineMetric {
}
fn preprocess(&self, vector: Vec<VectorElementType>) -> Vec<VectorElementType> {
let length: f32 = vector.iter().map(|x| x * x).sum();
let mut length: f32 = vector.iter().map(|x| x * x).sum();
length = length.sqrt();
let norm_vector = vector.iter().map(|x| x / length).collect();
return norm_vector;
}

88
lib/segment/src/spaces/tools.rs
View File

@@ -1,39 +1,61 @@
use crate::types::{Distance, Order, distance_order};
use serde::{Deserialize, Serialize};
use crate::types::Distance;
use std::collections::BinaryHeap;
use std::cmp::Reverse;
use crate::spaces::metric::Metric;
use crate::spaces::simple::{CosineMetric, DotProductMetric};
use crate::spaces::simple::{CosineMetric, DotProductMetric, EuclidMetric};
struct FixedLengthPriorityQueue<T> {
heap: BinaryHeap<T>,
/// This is a MinHeap by default - it will keep the largest elements, pop smallest
#[derive(Deserialize, Serialize, Clone, Debug)]
pub struct FixedLengthPriorityQueue<T: Ord> {
heap: BinaryHeap<Reverse<T>>,
length: usize,
}
impl<T: Ord> FixedLengthPriorityQueue<T> {
pub fn new(length: usize) -> Self {
assert!(length > 0);
FixedLengthPriorityQueue::<T> {
heap: BinaryHeap::new(),
heap: BinaryHeap::with_capacity(length + 1),
length,
}
}
pub fn push(&mut self, value: T) -> Option<T> {
self.heap.push(value);
return if self.heap.len() > self.length {
self.heap.pop()
if self.heap.len() < self.length {
self.heap.push(Reverse(value));
return None
}
return if self.heap.peek().unwrap().0 < value {
self.heap.push(Reverse(value));
self.heap.pop().map(|x| x.0)
} else {
None
};
Some(value)
}
}
pub fn into_vec(self) -> Vec<T> {
self.heap.into_sorted_vec()
self.heap.into_sorted_vec().into_iter().map(|x| x.0).collect()
}
pub fn into_iter(self) -> impl Iterator<Item=T> {
self.heap.into_sorted_vec().into_iter().map(|x| x.0)
}
pub fn iter(&self) -> impl Iterator<Item=&T> + '_ {
self.heap.iter().rev().map(|x| &x.0)
}
pub fn top(&self) -> Option<&T> { self.heap.peek().map(|x| &x.0) }
/// Return actual length of the queue
pub fn len(&self) -> usize {
self.heap.len()
}
}
pub fn peek_top_scores_iterable<I, E: Ord + Clone>(scores: I, top: usize, distance: &Distance) -> Vec<E>
pub fn peek_top_scores_iterable<I, E: Ord + Clone>(scores: I, top: usize) -> Vec<E>
where
I: Iterator<Item=E>,
{
@@ -41,40 +63,24 @@ pub fn peek_top_scores_iterable<I, E: Ord + Clone>(scores: I, top: usize, distan
return scores.collect();
}
let order = distance_order(&distance);
let res = match order {
Order::SmallBetter => {
// If small values is better - PQ should pop-out large values first.
// Hence is should be max-heap
let mut pq = FixedLengthPriorityQueue::new(top);
for score_point in scores {
pq.push(score_point.clone());
}
pq.into_vec()
}
Order::LargeBetter => {
let mut pq = FixedLengthPriorityQueue::new(top);
for score_point in scores {
pq.push(Reverse(score_point.clone()));
}
pq.into_vec()
.iter()
.map(|x| match x { Reverse(v) => v.clone() })
.collect()
}
};
return res;
// If big values is better - PQ should pop-out small values first.
// Hence is should be min-heap
let mut pq = FixedLengthPriorityQueue::new(top);
for score_point in scores {
pq.push(score_point.clone());
}
pq.into_vec()
}
pub fn peek_top_scores<E: Ord + Clone>(scores: &[E], top: usize, distance: &Distance) -> Vec<E> {
return peek_top_scores_iterable(scores.iter().cloned(), top, distance)
pub fn peek_top_scores<E: Ord + Clone>(scores: &[E], top: usize) -> Vec<E> {
return peek_top_scores_iterable(scores.iter().cloned(), top)
}
pub fn mertic_object(distance: &Distance) -> Box<dyn Metric> {
match distance {
Distance::Cosine => Box::new(CosineMetric {}),
Distance::Euclid => unimplemented!(),
Distance::Euclid => Box::new(EuclidMetric {}),
Distance::Dot => Box::new(DotProductMetric {}),
}
}
@@ -87,9 +93,7 @@ mod tests {
#[test]
fn test_peek_top() {
let data = vec![10, 20, 40, 5, 100, 33, 84, 65, 20, 43, 44, 42];
let res = peek_top_scores(&data, 3, &Distance::Dot);
let res = peek_top_scores(&data, 3);
assert_eq!(res, vec![100, 84, 65]);
let res = peek_top_scores(&data, 3, &Distance::Euclid);
assert_eq!(res, vec![5, 10, 20]);
}
}

54
lib/segment/src/types.rs
View File

@@ -6,7 +6,7 @@ use std::collections::{BTreeMap, HashSet, HashMap};
pub type PointIdType = u64;
/// Type of point index across all segments
pub type PointOffsetType = usize;
pub type PointOffsetType = u32;
/// Type of point index inside a segment
pub type PayloadKeyType = String;
pub type SeqNumberType = u64;
@@ -95,12 +95,10 @@ pub struct SegmentInfo {
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone, Copy, PartialEq)]
#[serde(rename_all = "snake_case")]
/// Additional parameters of the search
pub enum SearchParams {
pub struct SearchParams {
/// Params relevant to HNSW index
Hnsw {
/// Size of the beam in a beam-search. Larger the value - more accurate the result, more time required for search.
ef: usize
}
/// /// Size of the beam in a beam-search. Larger the value - more accurate the result, more time required for search.
pub hnsw_ef: Option<usize>
}
/// This function only stores mapping between distance and preferred result order
@@ -121,20 +119,29 @@ pub enum Indexes {
Plain {},
/// Use filterable HNSW index for approximate search. Is very fast even on a very huge collections,
/// but require additional space to store index and additional time to build it.
Hnsw {
/// Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.
m: usize,
/// Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.
ef_construct: usize,
},
Hnsw(HnswConfig),
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Copy, Clone, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub struct HnswConfig {
/// Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.
pub m: usize,
/// Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.
pub ef_construct: usize,
/// Minimal amount of points for additional payload-based indexing.
/// If payload chunk is smaller than `full_scan_threshold` additional indexing won't be used -
/// in this case full-scan search should be preferred by query planner and additional indexing is not required.
pub full_scan_threshold: usize
}
impl Default for HnswConfig {
fn default() -> Self { HnswConfig { m: 16, ef_construct: 100, full_scan_threshold: DEFAULT_FULL_SCAN_THRESHOLD }}
}
impl Indexes {
pub fn default_hnsw() -> Self {
Indexes::Hnsw {
m: 16,
ef_construct: 100,
}
Indexes::Hnsw(Default::default())
}
}
@@ -184,16 +191,19 @@ impl Default for StorageType {
pub struct SegmentConfig {
/// Size of a vectors used
pub vector_size: usize,
/// Type of distance function used for measuring distance between vectors
pub distance: Distance,
/// Type of index used for search
pub index: Indexes,
/// Payload Indexes
pub payload_index: Option<PayloadIndexType>,
/// Type of distance function used for measuring distance between vectors
pub distance: Distance,
/// Type of vector storage
pub storage_type: StorageType,
}
/// Default value based on https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md
pub const DEFAULT_FULL_SCAN_THRESHOLD: usize = 20_000;
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
#[serde(rename_all = "snake_case")]
pub struct SegmentState {
@@ -263,7 +273,7 @@ pub enum PayloadInterface {
KeywordShortcut(PayloadVariant<String>),
IntShortcut(PayloadVariant<i64>),
FloatShortcut(PayloadVariant<f64>),
Regular(PayloadInterfaceStrict),
Payload(PayloadInterfaceStrict),
}
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
@@ -279,7 +289,7 @@ pub enum PayloadInterfaceStrict {
// For tests
impl From<PayloadInterfaceStrict> for PayloadInterface {
fn from(x: PayloadInterfaceStrict) -> Self {
PayloadInterface::Regular(x)
PayloadInterface::Payload(x)
}
}
@@ -297,7 +307,7 @@ impl From<&PayloadInterfaceStrict> for PayloadType {
impl From<&PayloadInterface> for PayloadType {
fn from(interface: &PayloadInterface) -> Self {
match interface {
PayloadInterface::Regular(x) => x.into(),
PayloadInterface::Payload(x) => x.into(),
PayloadInterface::KeywordShortcut(x) => PayloadType::Keyword(x.to_list()),
PayloadInterface::FloatShortcut(x) => PayloadType::Float(x.to_list()),
PayloadInterface::IntShortcut(x) => PayloadType::Integer(x.to_list()),
@@ -550,7 +560,7 @@ mod tests {
#[test]
fn test_deny_unknown_fields() {
let query1 = r#"
let query1 = r#"
{
"wrong": "query"
}

356
lib/segment/src/vector_storage/memmap_vector_storage.rs
View File

@@ -1,169 +1,125 @@
use crate::vector_storage::vector_storage::{VectorStorage, ScoredPointOffset};
use crate::vector_storage::vector_storage::{VectorStorage, ScoredPointOffset, RawScorer};
use crate::entry::entry_point::OperationResult;
use std::ops::Range;
use std::path::{Path, PathBuf};
use std::fs::{File, OpenOptions, create_dir_all};
use memmap::{MmapOptions, Mmap, MmapMut};
use std::mem::{size_of, transmute};
use crate::types::{VectorElementType, PointOffsetType, Distance};
use std::fs::{OpenOptions, create_dir_all};
use std::mem::{size_of};
use crate::types::{VectorElementType, PointOffsetType, Distance, ScoreType};
use std::io::Write;
use crate::spaces::tools::{mertic_object, peek_top_scores};
use crate::common::error_logging::LogError;
use crate::spaces::tools::{mertic_object, peek_top_scores_iterable};
use crate::spaces::metric::Metric;
use crate::vector_storage::mmap_vectors::MmapVectors;
pub struct MemmapVectorStorage {
dim: usize,
num_vectors: usize,
mmap: Option<Mmap>,
deleted_mmap: Option<MmapMut>,
data_path: PathBuf,
deleted_path: PathBuf,
deleted_count: usize,
}
const HEADER_SIZE: usize = 4;
fn vf_to_u8<T>(v: &Vec<T>) -> &[u8] {
unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, v.len() * size_of::<T>()) }
}
pub struct MemmapRawScorer<'a> {
query: Vec<VectorElementType>,
metric: &'a Box<dyn Metric>,
mmap_store: &'a MmapVectors,
}
impl RawScorer for MemmapRawScorer<'_> {
fn score_points<'a>(&'a self, points: &'a mut dyn Iterator<Item=PointOffsetType>) -> Box<dyn Iterator<Item=ScoredPointOffset> + 'a> {
let res_iter = points
.filter(move |point| !self.mmap_store.deleted(*point).unwrap_or(true))
.map(move |point| {
let other_vector = self.mmap_store.raw_vector(point).unwrap();
ScoredPointOffset {
idx: point,
score: self.metric.similarity(&self.query, other_vector),
}
});
Box::new(res_iter)
}
fn check_point(&self, point: PointOffsetType) -> bool {
(point < self.mmap_store.num_vectors as PointOffsetType) && !self.mmap_store.deleted(point).unwrap_or(true)
}
fn score_point(&self, point: PointOffsetType) -> ScoreType {
let other_vector = self.mmap_store.raw_vector(point).unwrap();
self.metric.similarity(&self.query, other_vector)
}
fn score_internal(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> ScoreType {
let vector_a = self.mmap_store.raw_vector(point_a).unwrap();
let vector_b = self.mmap_store.raw_vector(point_b).unwrap();
return self.metric.similarity(vector_a, vector_b)
}
}
pub struct MemmapVectorStorage {
vectors_path: PathBuf,
deleted_path: PathBuf,
mmap_store: Option<MmapVectors>,
metric: Box<dyn Metric>
}
impl MemmapVectorStorage {
fn ensure_data_file_exists(path: &Path) -> OperationResult<()> {
if path.exists() {
return Ok(());
}
let mut file = File::create(path)?;
file.write(b"data")?;
Ok(())
}
fn ensure_deleted_file_exists(path: &Path) -> OperationResult<()> {
if path.exists() {
return Ok(());
}
let mut file = File::create(path)?;
file.write(b"drop")?;
Ok(())
}
fn open_read(path: &Path) -> OperationResult<Mmap> {
let file = OpenOptions::new()
.read(true)
.write(false)
.append(true)
.create(true)
.open(path)?;
let mmap = unsafe { MmapOptions::new().map(&file)? };
return Ok(mmap);
}
fn open_write(path: &Path) -> OperationResult<MmapMut> {
let file = OpenOptions::new()
.read(true)
.write(true)
.create(false)
.open(path)?;
let mmap = unsafe { MmapMut::map_mut(&file)? };
return Ok(mmap);
}
pub fn open(path: &Path, dim: usize) -> OperationResult<Self> {
pub fn open(path: &Path, dim: usize, distance: Distance) -> OperationResult<Self> {
create_dir_all(path)?;
let data_path = path.join("matrix.dat");
let vectors_path = path.join("matrix.dat");
let deleted_path = path.join("deleted.dat");
MemmapVectorStorage::ensure_data_file_exists(data_path.as_path()).describe("Create mmap data file")?;
MemmapVectorStorage::ensure_deleted_file_exists(deleted_path.as_path()).describe("Create mmap deleted flags file")?;
let mmap_store = MmapVectors::open(
vectors_path.as_path(),
deleted_path.as_path(),
dim,
)?;
let mmap = MemmapVectorStorage::open_read(&data_path).describe("Open mmap for reading")?;
let num_vectors = (mmap.len() - HEADER_SIZE) / dim / size_of::<VectorElementType>();
let deleted_mmap = MemmapVectorStorage::open_write(&deleted_path).describe("Open mmap for writing")?;
let deleted_count = (HEADER_SIZE..deleted_mmap.len())
.map(|idx| *deleted_mmap.get(idx).unwrap() as usize).sum();
let metric = mertic_object(&distance);
Ok(MemmapVectorStorage {
dim,
num_vectors,
mmap: Some(mmap),
deleted_mmap: Some(deleted_mmap),
data_path,
vectors_path,
deleted_path,
deleted_count,
mmap_store: Some(mmap_store),
metric
})
}
fn data_offset(&self, key: PointOffsetType) -> Option<usize> {
let vector_data_length = self.dim * size_of::<VectorElementType>();
let offset = key * vector_data_length + HEADER_SIZE;
if key >= self.num_vectors {
return None;
}
Some(offset)
}
fn raw_size(&self) -> usize {
self.dim * size_of::<VectorElementType>()
}
fn raw_vector_offset(&self, offset: usize) -> &[VectorElementType] {
let byte_slice = &self.mmap.as_ref().unwrap()[offset..(offset + self.raw_size())];
let arr: &[VectorElementType] = unsafe { transmute(byte_slice) };
return &arr[0..self.dim];
}
fn raw_vector(&self, key: PointOffsetType) -> Option<&[VectorElementType]> {
self.data_offset(key).map(|offset| self.raw_vector_offset(offset))
}
fn deleted(&self, key: PointOffsetType) -> Option<bool> {
self.deleted_mmap.as_ref().unwrap().get(HEADER_SIZE + key).map(|x| *x > 0)
}
}
impl VectorStorage for MemmapVectorStorage {
fn vector_dim(&self) -> usize {
self.dim
self.mmap_store.as_ref().unwrap().dim
}
fn vector_count(&self) -> usize {
self.num_vectors - self.deleted_count
self.mmap_store.as_ref().map(|store| store.num_vectors - store.deleted_count).unwrap()
}
fn deleted_count(&self) -> usize {
self.deleted_count
self.mmap_store.as_ref().unwrap().deleted_count
}
fn total_vector_count(&self) -> usize {
self.mmap_store.as_ref().unwrap().num_vectors
}
fn get_vector(&self, key: PointOffsetType) -> Option<Vec<VectorElementType>> {
match self.deleted(key) {
None => None,
Some(false) => self.data_offset(key).map(|offset| {
self.raw_vector_offset(offset).to_vec()
}),
Some(true) => None
}
self.mmap_store.as_ref().and_then(|x| x.get_vector(key))
}
fn put_vector(&mut self, _vector: &Vec<VectorElementType>) -> OperationResult<PointOffsetType> {
fn put_vector(&mut self, _vector: Vec<VectorElementType>) -> OperationResult<PointOffsetType> {
panic!("Can't put vector in mmap storage")
}
fn update_vector(&mut self, _key: usize, _vector: &Vec<VectorElementType>) -> OperationResult<usize> {
fn update_vector(&mut self, _key: PointOffsetType, _vector: Vec<VectorElementType>) -> OperationResult<PointOffsetType> {
panic!("Can't directly update vector in mmap storage")
}
fn update_from(&mut self, other: &dyn VectorStorage) -> OperationResult<Range<PointOffsetType>> {
self.mmap = None;
self.deleted_mmap = None;
let dim = self.vector_dim();
let start_index = self.num_vectors;
let mut end_index = self.num_vectors;
let start_index = self.mmap_store.as_ref().unwrap().num_vectors as PointOffsetType;
let mut end_index = start_index;
self.mmap_store = None;
{
let mut file = OpenOptions::new()
@@ -171,7 +127,7 @@ impl VectorStorage for MemmapVectorStorage {
.write(false)
.append(true)
.create(false)
.open(self.data_path.as_path())?;
.open(self.vectors_path.as_path())?;
for id in other.iter_ids() {
let vector = &other.get_vector(id).unwrap();
@@ -190,92 +146,103 @@ impl VectorStorage for MemmapVectorStorage {
.create(false)
.open(self.deleted_path.as_path())?;
let flags: Vec<u8> = vec![0; end_index - start_index];
let flags: Vec<u8> = vec![0; (end_index - start_index) as usize];
let flag_bytes = vf_to_u8(&flags);
file.write(flag_bytes)?;
file.flush()?;
}
let tmp_storage = Self::open(self.data_path.parent().unwrap(), self.dim)?;
self.mmap = tmp_storage.mmap;
self.deleted_mmap = tmp_storage.deleted_mmap;
self.num_vectors = tmp_storage.num_vectors;
self.deleted_count = tmp_storage.deleted_count;
self.mmap_store = Some(MmapVectors::open(
self.vectors_path.as_path(),
self.deleted_path.as_path(),
dim,
)?);
return Ok(start_index..end_index);
}
fn delete(&mut self, key: PointOffsetType) -> OperationResult<()> {
if key < self.num_vectors {
let mmap = self.deleted_mmap.as_mut().unwrap();
let flag = mmap.get_mut(key + HEADER_SIZE).unwrap();
self.mmap_store.as_mut().unwrap().delete(key)
}
if *flag == 0 {
*flag = 1;
self.deleted_count += 1;
}
}
Ok(())
fn is_deleted(&self, key: PointOffsetType) -> bool {
self.mmap_store.as_ref().unwrap().deleted(key).unwrap_or(false)
}
fn iter_ids(&self) -> Box<dyn Iterator<Item=PointOffsetType> + '_> {
let iter = (0..self.num_vectors)
.filter(move |id| !self.deleted(*id).unwrap());
let num_vectors = self.mmap_store.as_ref().unwrap().num_vectors;
let iter = (0..(num_vectors as PointOffsetType))
.filter(move |id| !self.mmap_store.as_ref().unwrap().deleted(*id).unwrap());
return Box::new(iter);
}
fn flush(&self) -> OperationResult<()> {
self.deleted_mmap.as_ref().unwrap().flush()?;
Ok(())
match self.mmap_store.as_ref() {
None => Ok(()),
Some(x) => x.flush()
}
}
fn raw_scorer(&self, vector: Vec<VectorElementType>) -> Box<dyn RawScorer + '_> {
Box::new(
MemmapRawScorer {
query: self.metric.preprocess(vector),
metric: &self.metric,
mmap_store: &self.mmap_store.as_ref().unwrap(),
}
)
}
fn raw_scorer_internal(&self, point_id: PointOffsetType) -> Box<dyn RawScorer + '_> {
Box::new(
MemmapRawScorer {
query: self.get_vector(point_id).unwrap(),
metric: &self.metric,
mmap_store: &self.mmap_store.as_ref().unwrap(),
}
)
}
fn score_points(
&self, vector: &Vec<VectorElementType>,
points: &[PointOffsetType],
points: &mut dyn Iterator<Item=PointOffsetType>,
top: usize,
distance: &Distance,
) -> Vec<ScoredPointOffset> {
let metric = mertic_object(distance);
let preprocessed_vector = metric.preprocess(vector.clone());
let scores: Vec<ScoredPointOffset> = points.iter()
.cloned()
.filter(|point| !self.deleted(*point).unwrap_or(true))
let preprocessed_vector = self.metric.preprocess(vector.clone());
let scores = points
.filter(|point| !self.mmap_store.as_ref().unwrap().deleted(*point).unwrap_or(true))
.map(|point| {
let other_vector =self.raw_vector(point).unwrap();
let other_vector = self.mmap_store.as_ref().unwrap().raw_vector(point).unwrap();
ScoredPointOffset {
idx: point,
score: metric.similarity(&preprocessed_vector, &other_vector),
score: self.metric.similarity(&preprocessed_vector, &other_vector),
}
}).collect();
return peek_top_scores(&scores, top, distance);
});
return peek_top_scores_iterable(scores, top);
}
fn score_all(&self, vector: &Vec<VectorElementType>, top: usize, distance: &Distance) -> Vec<ScoredPointOffset> {
let metric = mertic_object(distance);
let preprocessed_vector = metric.preprocess(vector.clone());
let scores: Vec<ScoredPointOffset> = self.iter_ids()
fn score_all(&self, vector: &Vec<VectorElementType>, top: usize) -> Vec<ScoredPointOffset> {
let preprocessed_vector = self.metric.preprocess(vector.clone());
let scores = self.iter_ids()
.map(|point| {
let other_vector = self.raw_vector(point).unwrap();
let other_vector = self.mmap_store.as_ref().unwrap().raw_vector(point).unwrap();
ScoredPointOffset {
idx: point,
score: metric.similarity(&preprocessed_vector, other_vector),
score: self.metric.similarity(&preprocessed_vector, other_vector),
}
}).collect();
});
return peek_top_scores(&scores, top, distance);
return peek_top_scores_iterable(scores, top);
}
fn score_internal(
&self,
point: PointOffsetType,
points: &[PointOffsetType],
points: &mut dyn Iterator<Item=PointOffsetType>,
top: usize,
distance: &Distance,
) -> Vec<ScoredPointOffset> {
let vector = self.get_vector(point).unwrap();
return self.score_points(&vector, points, top, distance);
return self.score_points(&vector, points, top);
}
}
@@ -285,11 +252,13 @@ mod tests {
use tempdir::TempDir;
use std::mem::transmute;
use crate::vector_storage::simple_vector_storage::SimpleVectorStorage;
use itertools::Itertools;
#[test]
fn test_basic_persistence() {
let dist = Distance::Dot;
let dir = TempDir::new("storage_dir").unwrap();
let mut storage = MemmapVectorStorage::open(dir.path(), 4).unwrap();
let mut storage = MemmapVectorStorage::open(dir.path(), 4, dist).unwrap();
let vec1 = vec![1.0, 0.0, 1.0, 1.0];
let vec2 = vec![1.0, 0.0, 1.0, 0.0];
@@ -299,11 +268,11 @@ mod tests {
{
let dir2 = TempDir::new("storage_dir2").unwrap();
let mut storage2 = SimpleVectorStorage::open(dir2.path(), 4).unwrap();
let mut storage2 = SimpleVectorStorage::open(dir2.path(), 4, dist).unwrap();
storage2.put_vector(&vec1).unwrap();
storage2.put_vector(&vec2).unwrap();
storage2.put_vector(&vec3).unwrap();
storage2.put_vector(vec1.clone()).unwrap();
storage2.put_vector(vec2.clone()).unwrap();
storage2.put_vector(vec3.clone()).unwrap();
storage.update_from(&storage2).unwrap();
}
@@ -319,9 +288,9 @@ mod tests {
{
let dir2 = TempDir::new("storage_dir2").unwrap();
let mut storage2 = SimpleVectorStorage::open(dir2.path(), 4).unwrap();
storage2.put_vector(&vec4).unwrap();
storage2.put_vector(&vec5).unwrap();
let mut storage2 = SimpleVectorStorage::open(dir2.path(), 4, dist).unwrap();
storage2.put_vector(vec4.clone()).unwrap();
storage2.put_vector(vec5.clone()).unwrap();
storage.update_from(&storage2).unwrap();
}
@@ -333,19 +302,58 @@ mod tests {
assert_eq!(stored_ids, vec![0, 1, 3, 4]);
let res = storage.score_all(&vec3, 2, &Distance::Dot);
let res = storage.score_all(&vec3, 2);
assert_eq!(res.len(), 2);
assert_ne!(res[0].idx, 2);
let res = storage.score_points(
&vec3, &vec![0, 1, 2, 3, 4], 2, &Distance::Dot);
&vec3, &mut vec![0, 1, 2, 3, 4].iter().cloned(), 2);
assert_eq!(res.len(), 2);
assert_ne!(res[0].idx, 2);
}
#[test]
fn test_mmap_raw_scorer() {
let dist = Distance::Dot;
let dir = TempDir::new("storage_dir").unwrap();
let mut storage = MemmapVectorStorage::open(dir.path(), 4, dist).unwrap();
let vec1 = vec![1.0, 0.0, 1.0, 1.0];
let vec2 = vec![1.0, 0.0, 1.0, 0.0];
let vec3 = vec![1.0, 1.0, 1.0, 1.0];
let vec4 = vec![1.0, 1.0, 0.0, 1.0];
let vec5 = vec![1.0, 0.0, 0.0, 0.0];
{
let dir2 = TempDir::new("storage_dir2").unwrap();
let mut storage2 = SimpleVectorStorage::open(dir2.path(), 4, dist).unwrap();
storage2.put_vector(vec1.clone()).unwrap();
storage2.put_vector(vec2.clone()).unwrap();
storage2.put_vector(vec3.clone()).unwrap();
storage2.put_vector(vec4.clone()).unwrap();
storage2.put_vector(vec5.clone()).unwrap();
storage.update_from(&storage2).unwrap();
}
let query = vec![-1.0, -1.0, -1.0, -1.0];
let query_points: Vec<PointOffsetType> = vec![0, 2, 4];
let scorer = storage.raw_scorer(query.clone());
let res = scorer.score_points(&mut query_points.iter().cloned()).collect_vec();
assert_eq!(res.len(), 3);
assert_eq!(res[0].idx, 0);
assert_eq!(res[1].idx, 2);
assert_eq!(res[2].idx, 4);
assert_eq!(res[2].score, -1.0);
}
#[test]
fn test_casts() {

131
lib/segment/src/vector_storage/mmap_vectors.rs Normal file
View File

@@ -0,0 +1,131 @@
use memmap::{Mmap, MmapMut, MmapOptions};
use crate::entry::entry_point::OperationResult;
use std::fs::{OpenOptions, File};
use std::path::Path;
use std::io::Write;
use std::mem::{size_of, transmute};
use crate::common::error_logging::LogError;
use crate::types::{VectorElementType, PointOffsetType};
const HEADER_SIZE: usize = 4;
const DELETED_HEADER: &[u8; 4] = b"drop";
const VECTORS_HEADER: &[u8; 4] = b"data";
pub struct MmapVectors {
pub dim: usize,
pub num_vectors: usize,
mmap: Mmap,
deleted_mmap: MmapMut,
pub deleted_count: usize,
}
fn open_read(path: &Path) -> OperationResult<Mmap> {
let file = OpenOptions::new()
.read(true)
.write(false)
.append(true)
.create(true)
.open(path)?;
let mmap = unsafe { MmapOptions::new().map(&file)? };
return Ok(mmap);
}
fn open_write(path: &Path) -> OperationResult<MmapMut> {
let file = OpenOptions::new()
.read(true)
.write(true)
.create(false)
.open(path)?;
let mmap = unsafe { MmapMut::map_mut(&file)? };
return Ok(mmap);
}
fn ensure_mmap_file_exists(path: &Path, header: &[u8]) -> OperationResult<()> {
if path.exists() {
return Ok(());
}
let mut file = File::create(path)?;
file.write(header)?;
Ok(())
}
impl MmapVectors {
pub fn open(vectors_path: &Path, deleted_path: &Path, dim: usize) -> OperationResult<Self> {
ensure_mmap_file_exists(vectors_path, VECTORS_HEADER).describe("Create mmap data file")?;
ensure_mmap_file_exists(deleted_path, DELETED_HEADER).describe("Create mmap deleted flags file")?;
let mmap = open_read(vectors_path).describe("Open mmap for reading")?;
let num_vectors = (mmap.len() - HEADER_SIZE) / dim / size_of::<VectorElementType>();
let deleted_mmap = open_write(deleted_path).describe("Open mmap for writing")?;
let deleted_count = (HEADER_SIZE..deleted_mmap.len())
.map(|idx| *deleted_mmap.get(idx).unwrap() as usize).sum();
Ok(MmapVectors {
dim,
num_vectors,
mmap,
deleted_mmap,
deleted_count,
})
}
pub fn data_offset(&self, key: PointOffsetType) -> Option<usize> {
let vector_data_length = self.dim * size_of::<VectorElementType>();
let offset = (key as usize) * vector_data_length + HEADER_SIZE;
if key >= (self.num_vectors as PointOffsetType) {
return None;
}
Some(offset)
}
pub fn raw_size(&self) -> usize {
self.dim * size_of::<VectorElementType>()
}
pub fn raw_vector_offset(&self, offset: usize) -> &[VectorElementType] {
let byte_slice = &self.mmap[offset..(offset + self.raw_size())];
let arr: &[VectorElementType] = unsafe { transmute(byte_slice) };
return &arr[0..self.dim];
}
pub fn raw_vector(&self, key: PointOffsetType) -> Option<&[VectorElementType]> {
self.data_offset(key).map(|offset| self.raw_vector_offset(offset))
}
pub fn deleted(&self, key: PointOffsetType) -> Option<bool> {
self.deleted_mmap.get(HEADER_SIZE + (key as usize)).map(|x| *x > 0)
}
/// Creates returns owned vector (copy of internal vector)
pub fn get_vector(&self, key: PointOffsetType) -> Option<Vec<VectorElementType>> {
match self.deleted(key) {
None => None,
Some(false) => self.data_offset(key).map(|offset| {
self.raw_vector_offset(offset).to_vec()
}),
Some(true) => None
}
}
pub fn delete(&mut self, key: PointOffsetType) -> OperationResult<()> {
if key < (self.num_vectors as PointOffsetType) {
let flag = self.deleted_mmap.get_mut((key as usize) + HEADER_SIZE).unwrap();
if *flag == 0 {
*flag = 1;
self.deleted_count += 1;
}
}
Ok(())
}
pub fn flush(&self) -> OperationResult<()> {
self.deleted_mmap.flush()?;
Ok(())
}
}

1
lib/segment/src/vector_storage/mod.rs
View File

@@ -1,3 +1,4 @@
pub mod vector_storage;
pub mod simple_vector_storage;
pub mod memmap_vector_storage;
mod mmap_vectors;

222
lib/segment/src/vector_storage/simple_vector_storage.rs
View File

@@ -1,4 +1,3 @@
use std::collections::HashSet;
use std::ops::Range;
use std::path::Path;
@@ -7,21 +6,25 @@ use rocksdb::{DB, IteratorMode, Options};
use serde::{Deserialize, Serialize};
use crate::entry::entry_point::OperationResult;
use crate::spaces::tools::{mertic_object, peek_top_scores};
use crate::types::{Distance, PointOffsetType, VectorElementType};
use crate::vector_storage::vector_storage::ScoredPointOffset;
use crate::spaces::tools::{mertic_object, peek_top_scores_iterable};
use crate::types::{Distance, PointOffsetType, VectorElementType, ScoreType};
use crate::vector_storage::vector_storage::{ScoredPointOffset, RawScorer};
use super::vector_storage::VectorStorage;
use std::mem::size_of;
use ndarray::{Array1, Array};
use crate::spaces::metric::Metric;
use bit_vec::BitVec;
/// Since sled is used for reading only during the initialization, large read cache is not required
const DB_CACHE_SIZE: usize = 10 * 1024 * 1024; // 10 mb
pub struct SimpleVectorStorage {
dim: usize,
metric: Box<dyn Metric>,
vectors: Vec<Array1<VectorElementType>>,
deleted: HashSet<PointOffsetType>,
deleted: BitVec,
deleted_count: usize,
store: DB,
}
@@ -32,11 +35,49 @@ struct StoredRecord {
pub vector: Vec<VectorElementType>,
}
pub struct SimpleRawScorer<'a> {
pub query: Array1<VectorElementType>,
pub metric: &'a Box<dyn Metric>,
pub vectors: &'a Vec<Array1<VectorElementType>>,
pub deleted: &'a BitVec,
}
impl RawScorer for SimpleRawScorer<'_> {
fn score_points<'a>(&'a self, points: &'a mut dyn Iterator<Item=PointOffsetType>) -> Box<dyn Iterator<Item=ScoredPointOffset> + 'a> {
let res_iter = points
.filter(move |point| !self.deleted[*point as usize])
.map(move |point| {
let other_vector = self.vectors.get(point as usize).unwrap();
ScoredPointOffset {
idx: point,
score: self.metric.blas_similarity(&self.query, other_vector),
}
});
Box::new(res_iter)
}
fn check_point(&self, point: PointOffsetType) -> bool {
(point < self.vectors.len() as PointOffsetType) && !self.deleted[point as usize]
}
fn score_point(&self, point: PointOffsetType) -> ScoreType {
let other_vector = &self.vectors[point as usize];
self.metric.blas_similarity(&self.query, other_vector)
}
fn score_internal(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> ScoreType {
let vector_a = &self.vectors[point_a as usize];
let vector_b = &self.vectors[point_b as usize];
return self.metric.blas_similarity(vector_a, vector_b);
}
}
impl SimpleVectorStorage {
pub fn open(path: &Path, dim: usize) -> OperationResult<Self> {
pub fn open(path: &Path, dim: usize, distance: Distance) -> OperationResult<Self> {
let mut vectors: Vec<Array1<VectorElementType>> = vec![];
let mut deleted: HashSet<PointOffsetType> = HashSet::new();
let mut deleted = BitVec::new();
let mut deleted_count = 0;
let mut options: Options = Options::default();
options.set_write_buffer_size(DB_CACHE_SIZE);
@@ -48,31 +89,41 @@ impl SimpleVectorStorage {
let point_id: PointOffsetType = bincode::deserialize(&key).unwrap();
let stored_record: StoredRecord = bincode::deserialize(&val).unwrap();
if stored_record.deleted {
deleted.insert(point_id);
deleted_count += 1;
}
if vectors.len() <= point_id {
vectors.resize(point_id + 1, Array::zeros(dim))
if vectors.len() <= (point_id as usize) {
vectors.resize((point_id + 1) as usize, Array::zeros(dim));
}
vectors[point_id].assign(&Array::from(stored_record.vector));
while deleted.len() <= (point_id as usize) {
deleted.push(false)
}
deleted.set(point_id as usize, stored_record.deleted);
vectors[point_id as usize].assign(&Array::from(stored_record.vector));
}
let metric = mertic_object(&distance);
debug!("Segment vectors: {}", vectors.len());
debug!("Estimated segment size {} MB", vectors.len() * dim * size_of::<VectorElementType>() / 1024 / 1024);
return Ok(SimpleVectorStorage {
dim,
metric,
vectors,
deleted,
deleted_count,
store,
});
}
fn update_stored(&self, point_id: PointOffsetType) -> OperationResult<()> {
let v = self.vectors.get(point_id).unwrap();
let v = self.vectors.get(point_id as usize).unwrap();
let record = StoredRecord {
deleted: self.deleted.contains(&point_id),
deleted: self.deleted[point_id as usize],
vector: v.to_vec(), // ToDo: try to reduce number of vector copies
};
self.store.put(
@@ -91,50 +142,69 @@ impl VectorStorage for SimpleVectorStorage {
}
fn vector_count(&self) -> usize {
self.vectors.len() - self.deleted.len()
self.vectors.len() - self.deleted_count
}
fn deleted_count(&self) -> usize {
return self.deleted.len();
self.deleted_count
}
fn total_vector_count(&self) -> usize {
self.vectors.len()
}
fn get_vector(&self, key: PointOffsetType) -> Option<Vec<VectorElementType>> {
if self.deleted.contains(&key) { return None; }
let vec = self.vectors.get(key)?.clone();
if self.deleted.get(key as usize).unwrap_or(true) { return None; }
let vec = self.vectors.get(key as usize)?.clone();
return Some(vec.to_vec());
}
fn put_vector(&mut self, vector: &Vec<VectorElementType>) -> OperationResult<PointOffsetType> {
fn put_vector(&mut self, vector: Vec<VectorElementType>) -> OperationResult<PointOffsetType> {
assert_eq!(self.dim, vector.len());
self.vectors.push(Array::from(vector.clone()));
self.update_stored(self.vectors.len() - 1)?;
return Ok(self.vectors.len() - 1);
self.vectors.push(Array::from(vector));
self.deleted.push(false);
let new_id = (self.vectors.len() - 1) as PointOffsetType;
self.update_stored(new_id)?;
return Ok(new_id);
}
fn update_vector(&mut self, key: usize, vector: &Vec<VectorElementType>) -> OperationResult<usize> {
self.vectors[key].assign(&Array::from(vector.clone()));
fn update_vector(&mut self, key: PointOffsetType, vector: Vec<VectorElementType>) -> OperationResult<PointOffsetType> {
self.vectors[key as usize].assign(&Array::from(vector));
self.update_stored(key)?;
return Ok(key);
}
fn update_from(&mut self, other: &dyn VectorStorage) -> OperationResult<Range<PointOffsetType>> {
let start_index = self.vectors.len();
let start_index = self.vectors.len() as PointOffsetType;
for id in other.iter_ids() {
self.put_vector(&other.get_vector(id).unwrap())?;
let other_vector = other.get_vector(id).unwrap();
// Do not perform preprocessing - vectors should be already processed
self.deleted.push(false);
self.vectors.push(Array::from(other_vector));
let new_id = (self.vectors.len() - 1) as PointOffsetType;
self.update_stored(new_id)?;
}
let end_index = self.vectors.len();
let end_index = self.vectors.len() as PointOffsetType;
return Ok(start_index..end_index);
}
fn delete(&mut self, key: PointOffsetType) -> OperationResult<()> {
self.deleted.insert(key);
if (key as usize) >= self.deleted.len() {
return Ok(())
}
if !self.deleted[key as usize] {
self.deleted_count += 1
}
self.deleted.set(key as usize, true);
self.update_stored(key)?;
Ok(())
}
fn iter_ids(&self) -> Box<dyn Iterator<Item=usize> + '_> {
let iter = (0..self.vectors.len())
.filter(move |id| !self.deleted.contains(id));
fn is_deleted(&self, key: PointOffsetType) -> bool { self.deleted[key as usize] }
fn iter_ids(&self) -> Box<dyn Iterator<Item=PointOffsetType> + '_> {
let iter = (0..self.vectors.len() as PointOffsetType)
.filter(move |id| !self.deleted[*id as usize]);
return Box::new(iter);
}
@@ -142,51 +212,64 @@ impl VectorStorage for SimpleVectorStorage {
Ok(self.store.flush()?)
}
fn raw_scorer(&self, vector: Vec<VectorElementType>) -> Box<dyn RawScorer + '_> {
Box::new(SimpleRawScorer {
query: Array::from(self.metric.preprocess(vector)),
metric: &self.metric,
vectors: &self.vectors,
deleted: &self.deleted,
})
}
fn raw_scorer_internal(&self, point_id: PointOffsetType) -> Box<dyn RawScorer + '_> {
Box::new(SimpleRawScorer {
query: self.vectors[point_id as usize].clone(),
metric: &self.metric,
vectors: &self.vectors,
deleted: &self.deleted,
})
}
fn score_points(
&self,
vector: &Vec<VectorElementType>,
points: &[PointOffsetType],
top: usize,
distance: &Distance,
points: &mut dyn Iterator<Item=PointOffsetType>,
top: usize
) -> Vec<ScoredPointOffset> {
let metric = mertic_object(distance);
let preprocessed_vector = Array::from(metric.preprocess(vector.clone()));
let scores: Vec<ScoredPointOffset> = points.iter()
.cloned()
.filter(|point| !self.deleted.contains(point))
let preprocessed_vector = Array::from(self.metric.preprocess(vector.clone()));
let scores = points
.filter(|point| !self.deleted[*point as usize])
.map(|point| {
let other_vector = self.vectors.get(point).unwrap();
let other_vector = self.vectors.get(point as usize).unwrap();
ScoredPointOffset {
idx: point,
score: metric.blas_similarity(&preprocessed_vector, other_vector),
score: self.metric.blas_similarity(&preprocessed_vector, other_vector),
}
}).collect();
return peek_top_scores(&scores, top, distance);
});
return peek_top_scores_iterable(scores, top);
}
fn score_all(&self, vector: &Vec<VectorElementType>, top: usize, distance: &Distance) -> Vec<ScoredPointOffset> {
let metric = mertic_object(distance);
let preprocessed_vector = Array::from(metric.preprocess(vector.clone()));
let scores: Vec<ScoredPointOffset> = self.vectors.iter()
fn score_all(&self, vector: &Vec<VectorElementType>, top: usize) -> Vec<ScoredPointOffset> {
let preprocessed_vector = Array::from(self.metric.preprocess(vector.clone()));
let scores = self.vectors.iter()
.enumerate()
.filter(|(point, _)| !self.deleted.contains(point))
.filter(|(point, _)| !self.deleted[*point])
.map(|(point, other_vector)| ScoredPointOffset {
idx: point,
score: metric.blas_similarity(&preprocessed_vector, other_vector),
}).collect();
return peek_top_scores(&scores, top, distance);
idx: point as PointOffsetType,
score: self.metric.blas_similarity(&preprocessed_vector, other_vector),
});
return peek_top_scores_iterable(scores, top);
}
fn score_internal(
&self,
point: PointOffsetType,
points: &[PointOffsetType],
points: &mut dyn Iterator<Item=PointOffsetType>,
top: usize,
distance: &Distance,
) -> Vec<ScoredPointOffset> {
let vector = self.get_vector(point).unwrap();
return self.score_points(&vector, points, top, distance);
return self.score_points(&vector, points, top);
}
}
@@ -196,24 +279,25 @@ mod tests {
use tempdir::TempDir;
use super::*;
use itertools::Itertools;
#[test]
fn test_score_points() {
let dir = TempDir::new("storage_dir").unwrap();
let distance = Distance::Dot;
let dim = 4;
let mut storage = SimpleVectorStorage::open(dir.path(), dim).unwrap();
let mut storage = SimpleVectorStorage::open(dir.path(), dim, distance).unwrap();
let vec0 = vec![1.0, 0.0, 1.0, 1.0];
let vec1 = vec![1.0, 0.0, 1.0, 0.0];
let vec2 = vec![1.0, 1.0, 1.0, 1.0];
let vec3 = vec![1.0, 1.0, 0.0, 1.0];
let vec4 = vec![1.0, 0.0, 0.0, 0.0];
let _id1 = storage.put_vector(&vec0).unwrap();
let id2 = storage.put_vector(&vec1).unwrap();
let _id3 = storage.put_vector(&vec2).unwrap();
let _id4 = storage.put_vector(&vec3).unwrap();
let id5 = storage.put_vector(&vec4).unwrap();
let _id1 = storage.put_vector(vec0.clone()).unwrap();
let id2 = storage.put_vector(vec1.clone()).unwrap();
let _id3 = storage.put_vector(vec2.clone()).unwrap();
let _id4 = storage.put_vector(vec3.clone()).unwrap();
let id5 = storage.put_vector(vec4.clone()).unwrap();
assert_eq!(id2, 1);
assert_eq!(id5, 4);
@@ -222,9 +306,8 @@ mod tests {
let closest = storage.score_points(
&query,
&[0, 1, 2, 3, 4],
&mut [0, 1, 2, 3, 4].iter().cloned(),
2,
&distance,
);
let top_idx = match closest.get(0) {
@@ -242,15 +325,26 @@ mod tests {
let closest = storage.score_points(
&query,
&[0, 1, 2, 3, 4],
&mut [0, 1, 2, 3, 4].iter().cloned(),
2,
&distance,
);
let raw_scorer = storage.raw_scorer(query.clone());
let query_points = vec![0, 1, 2, 3, 4];
let mut query_points1 = query_points.iter().cloned();
let mut query_points2 = query_points.iter().cloned();
let raw_res1 = raw_scorer.score_points(&mut query_points1).collect_vec();
let raw_res2 = raw_scorer.score_points(&mut query_points2).collect_vec();
assert_eq!(raw_res1, raw_res2);
let _top_idx = match closest.get(0) {
Some(scored_point) => {
assert_ne!(scored_point.idx, 2);
assert_eq!(&raw_res1[scored_point.idx as usize], scored_point);
}
None => { assert!(false, "No close vector found!") }
};

55
lib/segment/src/vector_storage/vector_storage.rs
View File

@@ -1,8 +1,9 @@
use crate::types::{PointOffsetType, ScoreType, VectorElementType, Distance};
use crate::types::{PointOffsetType, ScoreType, VectorElementType};
use std::cmp::{Ordering};
use ordered_float::OrderedFloat;
use crate::entry::entry_point::OperationResult;
use std::ops::Range;
use rand::Rng;
#[derive(Copy, Clone, PartialEq, Debug)]
@@ -26,6 +27,21 @@ impl PartialOrd for ScoredPointOffset {
}
/// Optimized scorer for multiple scoring requests comparing with a single query
/// Holds current query and params, receives only subset of points to score
pub trait RawScorer {
// ToDo: Replace boxed iterator with callback and make a benchmark (-4% on benchmarks, but ugly)
fn score_points<'a>(&'a self, points: &'a mut dyn Iterator<Item=PointOffsetType>) -> Box<dyn Iterator<Item=ScoredPointOffset> + 'a>;
/// Return true if point satisfies current search context (exists and not deleted)
fn check_point(&self, point: PointOffsetType) -> bool;
/// Score stored vector with vector under the given index
fn score_point(&self, point: PointOffsetType) -> ScoreType;
/// Return distance between stored points selected by ids
/// Panics if any id is out of range
fn score_internal(&self, point_a: PointOffsetType, point_b: PointOffsetType) -> ScoreType;
}
/// Trait for vector storage
/// El - type of vector element, expected numerical type
/// Storage operates with internal IDs (PointOffsetType), which always starts with zero and have no skips
@@ -33,34 +49,49 @@ pub trait VectorStorage {
fn vector_dim(&self) -> usize;
fn vector_count(&self) -> usize; /// Number of searchable vectors (not deleted)
fn deleted_count(&self) -> usize; /// Number of vectors, marked as deleted but still stored
fn total_vector_count(&self) -> usize; /// Number of all stored vectors including deleted
fn get_vector(&self, key: PointOffsetType) -> Option<Vec<VectorElementType>>;
fn put_vector(&mut self, vector: &Vec<VectorElementType>) -> OperationResult<PointOffsetType>;
fn update_vector(&mut self, key: PointOffsetType, vector: &Vec<VectorElementType>) -> OperationResult<PointOffsetType>;
fn put_vector(&mut self, vector: Vec<VectorElementType>) -> OperationResult<PointOffsetType>;
fn update_vector(&mut self, key: PointOffsetType, vector: Vec<VectorElementType>) -> OperationResult<PointOffsetType>;
fn update_from(&mut self, other: &dyn VectorStorage) -> OperationResult<Range<PointOffsetType>>;
fn delete(&mut self, key: PointOffsetType) -> OperationResult<()>;
fn iter_ids(&self) -> Box<dyn Iterator<Item=PointOffsetType> + '_>;
fn is_deleted(&self, key: PointOffsetType) -> bool;
fn iter_ids(&self) -> Box<dyn Iterator<Item=PointOffsetType> + '_>; /// Iterator over not-deleted ids
fn flush(&self) -> OperationResult<()>;
/// Generate a RawScorer object which contains all required context for searching similar vector
fn raw_scorer(&self, vector: Vec<VectorElementType>) -> Box<dyn RawScorer + '_>;
/// Same as `raw_scorer` but uses internal vector for search, avoids double pre-processing
fn raw_scorer_internal(&self, point_id: PointOffsetType) -> Box<dyn RawScorer + '_>;
fn score_points(
&self,
vector: &Vec<VectorElementType>,
points: &[PointOffsetType],
top: usize,
distance: &Distance,
points: &mut dyn Iterator<Item=PointOffsetType>,
top: usize
) -> Vec<ScoredPointOffset>;
fn score_all(
&self,
vector: &Vec<VectorElementType>,
top: usize,
distance: &Distance
top: usize
) -> Vec<ScoredPointOffset>;
fn score_internal(
&self,
point: PointOffsetType,
points: &[PointOffsetType],
top: usize,
distance: &Distance
points: &mut dyn Iterator<Item=PointOffsetType>,
top: usize
) -> Vec<ScoredPointOffset>;
/// Iterator over `n` random ids which are not deleted
fn sample_ids(&self) -> Box<dyn Iterator<Item=PointOffsetType> + '_> {
let total = self.total_vector_count() as PointOffsetType;
let mut rng = rand::thread_rng();
Box::new((0..total)
.map(move |_| rng.gen_range(0..total))
.filter(move |x| !self.is_deleted(*x))
)
}
}

130
lib/segment/tests/filtrable_hnsw_test.rs Normal file
View File

@@ -0,0 +1,130 @@
#[cfg(test)]
mod tests {
use tempdir::TempDir;
use segment::types::{StorageType, Distance, PayloadIndexType, Indexes, SegmentConfig, TheMap, PayloadKeyType, PayloadType, SeqNumberType, PointIdType, Condition, FieldCondition, Filter, Range, SearchParams, HnswConfig};
use segment::segment_constructor::segment_constructor::build_segment;
use segment::fixtures::payload_fixtures::{random_vector, random_int_payload};
use segment::entry::entry_point::SegmentEntry;
use segment::index::struct_payload_index::StructPayloadIndex;
use segment::index::index::{PayloadIndex, VectorIndex};
use segment::index::hnsw_index::hnsw::HNSWIndex;
use std::sync::Arc;
use atomic_refcell::AtomicRefCell;
use itertools::Itertools;
use rand::{thread_rng, Rng};
#[test]
fn test_filterable_hnsw() {
let dim = 8;
let m = 8;
let num_vectors: PointIdType = 5_000;
let ef = 32;
let ef_construct = 16;
let distance = Distance::Cosine;
let indexing_threshold = 500;
let num_payload_values = 2;
let mut rnd = thread_rng();
let dir = TempDir::new("segment_dir").unwrap();
let payload_index_dir = TempDir::new("payload_index_dir").unwrap();
let hnsw_dir = TempDir::new("hnsw_dir").unwrap();
let config = SegmentConfig {
vector_size: dim,
index: Indexes::Plain {},
payload_index: Some(PayloadIndexType::Plain),
storage_type: StorageType::InMemory,
distance,
};
let int_key = "int".to_string();
let mut segment = build_segment(dir.path(), &config).unwrap();
for idx in 0..num_vectors {
let vector = random_vector(&mut rnd, dim);
let mut payload: TheMap<PayloadKeyType, PayloadType> = Default::default();
payload.insert(int_key.clone(), random_int_payload(&mut rnd, num_payload_values));
segment.upsert_point(idx as SeqNumberType, idx, &vector).unwrap();
segment.set_full_payload(idx as SeqNumberType, idx, payload.clone()).unwrap();
}
// let opnum = num_vectors + 1;
let payload_index = StructPayloadIndex::open(
segment.condition_checker.clone(),
segment.vector_storage.clone(),
segment.payload_storage.clone(),
segment.id_mapper.clone(),
payload_index_dir.path(),
).unwrap();
let payload_index_ptr = Arc::new(AtomicRefCell::new(payload_index));
let hnsw_config = HnswConfig {
m,
ef_construct,
full_scan_threshold: indexing_threshold
};
let mut hnsw_index = HNSWIndex::open(
hnsw_dir.path(),
segment.condition_checker.clone(),
segment.vector_storage.clone(),
payload_index_ptr.clone(),
hnsw_config
).unwrap();
hnsw_index.build_index().unwrap();
payload_index_ptr.borrow_mut().set_indexed(&int_key).unwrap();
let borrowed_payload_index = payload_index_ptr.borrow();
let blocks = borrowed_payload_index.payload_blocks(indexing_threshold).collect_vec();
assert_eq!(blocks.len(), num_vectors as usize / indexing_threshold * 2);
hnsw_index.build_index().unwrap();
let top = 3;
let mut hits = 0;
let attempts = 100;
for _i in 0..attempts {
let query = random_vector(&mut rnd, dim);
let range_size = 40;
let left_range = rnd.gen_range(0..400);
let right_range = left_range + range_size;
let filter = Filter::new_must(Condition::Field(FieldCondition {
key: int_key.clone(),
r#match: None,
range: Some(Range {
lt: None,
gt: None,
gte: Some(left_range as f64),
lte: Some(right_range as f64),
}),
geo_bounding_box: None,
geo_radius: None,
}));
let filter_query = Some(&filter);
// let filter_query = None;
let index_result = hnsw_index.search_with_graph(
&query,
filter_query,
top,
Some(&SearchParams { hnsw_ef: Some(ef) })
);
let plain_result = segment.vector_index.borrow().search(&query, filter_query, top, None);
if plain_result == index_result {
hits += 1;
}
}
assert!(attempts - hits < 5); // Not more than 5% failures
eprintln!("hits = {:#?} out of {}", hits, attempts);
}
}

170
lib/segment/tests/payload_index_test.rs
View File

@@ -1,122 +1,81 @@
#[cfg(test)]
mod tests {
use rand::prelude::ThreadRng;
use rand::seq::SliceRandom;
use segment::types::{PayloadType, VectorElementType, SegmentConfig, Indexes, PayloadIndexType, Distance, StorageType, TheMap, PayloadKeyType, Filter, Condition, FieldCondition, Match, Range as RangeConditionl};
use rand::Rng;
use segment::fixtures::payload_fixtures::{random_vector, random_keyword_payload, random_int_payload, random_filter};
use tempdir::TempDir;
use segment::types::{SegmentConfig, Indexes, PayloadIndexType, StorageType, Distance, TheMap, PayloadKeyType, PayloadType, Filter, Condition, FieldCondition, Range};
use segment::segment_constructor::segment_constructor::build_segment;
use segment::entry::entry_point::SegmentEntry;
use itertools::Itertools;
use std::ops::Range;
const ADJECTIVE: &'static [&'static str] = &[
"jobless",
"rightful",
"breakable",
"impartial",
"shocking",
"faded",
"phobic",
"overt",
"like",
"wide-eyed",
"broad",
];
#[test]
fn test_cardinality_estimation() {
let mut rnd = rand::thread_rng();
const NOUN: &'static [&'static str] = &[
"territory",
"jam",
"neck",
"chicken",
"cap",
"kiss",
"veil",
"trail",
"size",
"digestion",
"rod",
"seed",
];
let dir1 = TempDir::new("segment1_dir").unwrap();
let dim = 5;
const INT_RANGE: Range<i64> = 0..500;
let
config = SegmentConfig {
vector_size: dim,
index: Indexes::Plain {},
payload_index: Some(PayloadIndexType::Struct),
storage_type: StorageType::InMemory,
distance: Distance::Dot,
};
fn random_keyword(rnd_gen: &mut ThreadRng) -> String {
let random_adj = ADJECTIVE.choose(rnd_gen).unwrap();
let random_noun = NOUN.choose(rnd_gen).unwrap();
format!("{} {}", random_adj, random_noun)
}
let str_key = "kvd".to_string();
let int_key = "int".to_string();
fn random_keyword_payload(rnd_gen: &mut ThreadRng) -> PayloadType {
PayloadType::Keyword(vec![random_keyword(rnd_gen)])
}
let num_points = 10000;
let mut struct_segment = build_segment(dir1.path(), &config).unwrap();
fn random_int_payload(rnd_gen: &mut ThreadRng) -> PayloadType {
let val1: i64 = rnd_gen.gen_range(INT_RANGE);
let val2: i64 = rnd_gen.gen_range(INT_RANGE);
PayloadType::Integer(vec![val1, val2])
}
let mut opnum = 0;
for idx in 0..num_points {
let vector = random_vector(&mut rnd, dim);
let mut payload: TheMap<PayloadKeyType, PayloadType> = Default::default();
payload.insert(str_key.clone(), random_keyword_payload(&mut rnd));
payload.insert(int_key.clone(), random_int_payload(&mut rnd, 2));
fn random_vector(rnd_gen: &mut ThreadRng, size: usize) -> Vec<VectorElementType> {
(0..size).map(|_| rnd_gen.gen()).collect()
}
struct_segment.upsert_point(opnum, idx, &vector).unwrap();
struct_segment.set_full_payload(opnum, idx, payload.clone()).unwrap();
fn random_field_condition(rnd_gen: &mut ThreadRng) -> Condition {
let kv_or_int: bool = rnd_gen.gen();
match kv_or_int {
true => Condition::Field(FieldCondition {
key: "kvd".to_string(),
r#match: Some(Match {
keyword: Some(random_keyword(rnd_gen)),
integer: None,
}),
range: None,
geo_bounding_box: None,
geo_radius: None,
opnum += 1;
}
struct_segment.create_field_index(opnum, &str_key).unwrap();
struct_segment.create_field_index(opnum, &int_key).unwrap();
let filter = Filter::new_must(Condition::Field(FieldCondition {
key: int_key,
r#match: None,
range: Some(Range {
lt: None,
gt: None,
gte: Some(50.),
lte: Some(100.),
}),
false => Condition::Field(FieldCondition {
key: "int".to_string(),
r#match: None,
range: Some(RangeConditionl {
lt: None,
gt: None,
gte: Some(rnd_gen.gen_range(INT_RANGE) as f64),
lte: Some(rnd_gen.gen_range(INT_RANGE) as f64),
}),
geo_bounding_box: None,
geo_radius: None,
})
}
}
geo_bounding_box: None,
geo_radius: None,
}));
fn random_filter(rnd_gen: &mut ThreadRng) -> Filter {
let mut rnd1 = rand::thread_rng();
let estimation = struct_segment.payload_index
.borrow()
.estimate_cardinality(&filter);
let should_conditions = (0..=2)
.take_while(|_| rnd1.gen::<f64>() > 0.6)
.map(|_| random_field_condition(rnd_gen))
.collect_vec();
let checker = struct_segment.condition_checker.borrow();
let should_conditions_opt = match should_conditions.is_empty() {
false => Some(should_conditions),
true => None,
};
let exact = struct_segment.vector_storage
.borrow()
.iter_ids()
.filter(|x| checker.check(*x, &filter))
.collect_vec()
.len();
let must_conditions = (0..=2)
.take_while(|_| rnd1.gen::<f64>() > 0.6)
.map(|_| random_field_condition(rnd_gen))
.collect_vec();
eprintln!("exact = {:#?}", exact);
eprintln!("estimation = {:#?}", estimation);
let must_conditions_opt = match must_conditions.is_empty() {
false => Some(must_conditions),
true => None,
};
Filter {
should: should_conditions_opt,
must: must_conditions_opt,
must_not: None,
}
assert!(exact <= estimation.max);
assert!(exact >= estimation.min);
}
#[test]
@@ -145,13 +104,14 @@ mod tests {
let int_key = "int".to_string();
let num_points = 1000;
let num_int_values = 2;
let mut opnum = 0;
for idx in 0..num_points {
let vector = random_vector(&mut rnd, dim);
let mut payload: TheMap<PayloadKeyType, PayloadType> = Default::default();
payload.insert(str_key.clone(), random_keyword_payload(&mut rnd));
payload.insert(int_key.clone(), random_int_payload(&mut rnd));
payload.insert(int_key.clone(), random_int_payload(&mut rnd, num_int_values));
plain_segment.upsert_point(idx, idx, &vector).unwrap();
struct_segment.upsert_point(idx, idx, &vector).unwrap();
@@ -165,8 +125,8 @@ mod tests {
struct_segment.create_field_index(opnum, &str_key).unwrap();
struct_segment.create_field_index(opnum, &int_key).unwrap();
for _i in 0..100 {
let attempts = 100;
for _i in 0..attempts {
let query_vector = random_vector(&mut rnd, dim);
let query_filter = random_filter(&mut rnd);
@@ -175,9 +135,9 @@ mod tests {
let estimation = struct_segment.payload_index.borrow().estimate_cardinality(&query_filter);
assert!(estimation.min <= estimation.exp);
assert!(estimation.exp <= estimation.max);
assert!(estimation.max <= num_points as usize);
assert!(estimation.min <= estimation.exp, "{:#?}", estimation);
assert!(estimation.exp <= estimation.max, "{:#?}", estimation);
assert!(estimation.max <= num_points as usize, "{:#?}", estimation);
plain_result
.iter()

2
lib/storage/src/content_manager/errors.rs
View File

@@ -1,5 +1,5 @@
use thiserror::Error;
use collection::collection::CollectionError;
use collection::operations::types::CollectionError;
use sled::Error;
use sled::transaction::TransactionError;
use std::io::Error as IoError;

18
lib/storage/src/content_manager/storage_ops.rs
View File

@@ -1,6 +1,7 @@
use serde::{Deserialize, Serialize};
use schemars::{JsonSchema};
use segment::types::{Distance, Indexes};
use segment::types::{Distance};
use collection::operations::config_diff::{HnswConfigDiff, WalConfigDiff, OptimizersConfigDiff};
#[derive(Debug, Deserialize, Serialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
@@ -30,7 +31,20 @@ pub enum StorageOperations {
name: String,
vector_size: usize,
distance: Distance,
index: Option<Indexes>,
/// Custom params for HNSW index. If none - values from service configuration file are used.
hnsw_config: Option<HnswConfigDiff>,
/// Custom params for WAL. If none - values from service configuration file are used.
wal_config: Option<WalConfigDiff>,
/// Custom params for Optimizers. If none - values from service configuration file are used.
optimizers_config: Option<OptimizersConfigDiff>
},
/// Update parameters of the existing collection
UpdateCollection {
name: String,
/// Custom params for Optimizers. If none - values from service configuration file are used.
/// This operation is blocking, it will only proceed ones all current optimizations are complete
optimizers_config: Option<OptimizersConfigDiff>
// ToDo: Allow updates for other configuration params as well
},
/// Delete collection with given name
DeleteCollection(String),

65
lib/storage/src/content_manager/toc.rs
View File

@@ -11,16 +11,16 @@ use sled::{Config, Db};
use sled::transaction::UnabortableTransactionError;
use tokio::runtime;
use tokio::runtime::Runtime;
use wal::WalOptions;
use collection::collection::Collection;
use collection::collection_builder::collection_builder::build_collection;
use collection::collection_builder::collection_loader::load_collection;
use segment::types::SegmentConfig;
use crate::content_manager::errors::StorageError;
use crate::content_manager::storage_ops::{AliasOperations, StorageOperations};
use crate::types::StorageConfig;
use collection::config::CollectionParams;
use collection::operations::config_diff::{DiffConfig};
/// Since sled is used for reading only during the initialization, large read cache is not required
const SLED_CACHE_SIZE: u64 = 1 * 1024 * 1024; // 1 mb
@@ -59,16 +59,10 @@ impl TableOfContent {
for entry in collection_paths {
let collection_path = entry.unwrap().path();
let collection_name = collection_path.file_name().unwrap().to_str().unwrap().to_string();
let wal_options = WalOptions {
segment_capacity: storage_config.wal.wal_capacity_mb * 1024 * 1024,
segment_queue_len: storage_config.wal.wal_segments_ahead,
};
let collection = load_collection(
collection_path.as_path(),
&wal_options,
search_runtime.clone(),
&storage_config.optimizers,
);
collections.insert(collection_name, Arc::new(collection));
@@ -151,36 +145,56 @@ impl TableOfContent {
name: collection_name,
vector_size,
distance,
index
hnsw_config: hnsw_config_diff,
wal_config: wal_config_diff,
optimizers_config: optimizers_config_diff,
} => {
self.validate_collection_not_exists(&collection_name)?;
let wal_options = WalOptions {
segment_capacity: self.storage_config.wal.wal_capacity_mb * 1024 * 1024,
segment_queue_len: self.storage_config.wal.wal_segments_ahead,
};
let collection_path = self.create_collection_path(&collection_name)?;
let segment_config = SegmentConfig {
let collection_params = CollectionParams {
vector_size,
index: index.unwrap_or(Default::default()),
payload_index: Some(Default::default()),
distance,
storage_type: Default::default(),
};
let wal_config = match wal_config_diff {
None => self.storage_config.wal.clone(),
Some(diff) => diff.update(&self.storage_config.wal)?
};
let segment = build_collection(
let optimizers_config = match optimizers_config_diff {
None => self.storage_config.optimizers.clone(),
Some(diff) => diff.update(&self.storage_config.optimizers)?,
};
let hnsw_config = match hnsw_config_diff {
None => self.storage_config.hnsw_index.clone(),
Some(diff) => diff.update(&self.storage_config.hnsw_index)?
};
let collection = build_collection(
Path::new(&collection_path),
&wal_options,
&segment_config,
&wal_config,
&collection_params,
self.search_runtime.clone(),
&self.storage_config.optimizers,
&optimizers_config,
&hnsw_config,
)?;
let mut write_collections = self.collections.write();
write_collections.insert(collection_name, Arc::new(segment));
write_collections.insert(collection_name, Arc::new(collection));
Ok(true)
}
StorageOperations::UpdateCollection {
name,
optimizers_config
} => {
let collection = self.get_collection(&name)?;
match optimizers_config {
None => {}
Some(new_optimizers_config) => {
collection.update_optimizer_params(new_optimizers_config)?
}
}
Ok(true)
}
StorageOperations::DeleteCollection(collection_name) => {
@@ -233,6 +247,7 @@ impl TableOfContent {
pub fn get_collection(&self, collection_name: &str) -> Result<Arc<Collection>, StorageError> {
let read_collection = self.collections.read();
let real_collection_name = self.resolve_name(collection_name)?;
// resolve_name already checked collection existence, unwrap is safe here
Ok(read_collection.get(&real_collection_name).unwrap().clone())
}

10
lib/storage/src/types.rs
View File

@@ -1,6 +1,8 @@
use serde::{Deserialize, Serialize};
use schemars::{JsonSchema};
use collection::collection_builder::optimizers_builder::OptimizersConfig;
use collection::config::WalConfig;
use segment::types::HnswConfig;
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
@@ -9,18 +11,12 @@ pub struct PerformanceConfig {
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
pub struct WalConfig {
pub wal_capacity_mb: usize,
pub wal_segments_ahead: usize,
}
#[derive(Debug, Deserialize, Serialize, JsonSchema, Clone)]
pub struct StorageConfig {
pub storage_path: String,
pub optimizers: OptimizersConfig,
pub wal: WalConfig,
pub performance: PerformanceConfig,
pub hnsw_index: HnswConfig
}

562
openapi/models.json
View File

@@ -25,7 +25,8 @@
}
}
}
}
},
"additionalProperties": false
},
{
"description": "Delete alias if exists",
@@ -45,7 +46,8 @@
}
}
}
}
},
"additionalProperties": false
},
{
"description": "Change alias to a new one",
@@ -69,10 +71,34 @@
}
}
}
}
},
"additionalProperties": false
}
]
},
"CollectionConfig": {
"type": "object",
"required": [
"hnsw_config",
"optimizer_config",
"params",
"wal_config"
],
"properties": {
"hnsw_config": {
"$ref": "#/components/schemas/HnswConfig"
},
"optimizer_config": {
"$ref": "#/components/schemas/OptimizersConfig"
},
"params": {
"$ref": "#/components/schemas/CollectionParams"
},
"wal_config": {
"$ref": "#/components/schemas/WalConfig"
}
}
},
"CollectionDescription": {
"type": "object",
"required": [
@@ -92,11 +118,12 @@
"disk_data_size",
"ram_data_size",
"segments_count",
"status",
"vectors_count"
],
"properties": {
"config": {
"$ref": "#/components/schemas/SegmentConfig"
"$ref": "#/components/schemas/CollectionConfig"
},
"disk_data_size": {
"description": "Disk space, used by collection",
@@ -116,6 +143,9 @@
"format": "uint",
"minimum": 0
},
"status": {
"$ref": "#/components/schemas/CollectionStatus"
},
"vectors_count": {
"description": "Number of vectors in collection",
"type": "integer",
@@ -124,6 +154,32 @@
}
}
},
"CollectionParams": {
"type": "object",
"required": [
"distance",
"vector_size"
],
"properties": {
"distance": {
"$ref": "#/components/schemas/Distance"
},
"vector_size": {
"description": "Size of a vectors used",
"type": "integer",
"format": "uint",
"minimum": 0
}
}
},
"CollectionStatus": {
"type": "string",
"enum": [
"green",
"yellow",
"red"
]
},
"CollectionUpdateOperations": {
"anyOf": [
{
@@ -240,7 +296,8 @@
"create_index": {
"type": "string"
}
}
},
"additionalProperties": false
},
{
"description": "Delete index for the field",
@@ -252,7 +309,8 @@
"delete_index": {
"type": "string"
}
}
},
"additionalProperties": false
}
]
},
@@ -347,69 +405,64 @@
"type": "integer",
"format": "uint64",
"minimum": 0
}
},
"uniqueItems": true
}
}
},
"Indexes": {
"anyOf": [
{
"description": "Do not use any index, scan whole vector collection during search. Guarantee 100% precision, but may be time consuming on large collections.",
"type": "object",
"required": [
"options",
"type"
],
"properties": {
"options": {
"type": "object"
},
"type": {
"type": "string",
"enum": [
"plain"
]
}
}
"HnswConfig": {
"type": "object",
"required": [
"ef_construct",
"full_scan_threshold",
"m"
],
"properties": {
"ef_construct": {
"description": "Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.",
"type": "integer",
"format": "uint",
"minimum": 0
},
{
"description": "Use filterable HNSW index for approximate search. Is very fast even on a very huge collections, but require additional space to store index and additional time to build it.",
"type": "object",
"required": [
"options",
"type"
],
"properties": {
"options": {
"type": "object",
"required": [
"ef_construct",
"m"
],
"properties": {
"ef_construct": {
"description": "Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.",
"type": "integer",
"format": "uint",
"minimum": 0
},
"m": {
"description": "Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.",
"type": "integer",
"format": "uint",
"minimum": 0
}
}
},
"type": {
"type": "string",
"enum": [
"hnsw"
]
}
}
"full_scan_threshold": {
"description": "Minimal amount of points for additional payload-based indexing. If payload chunk is smaller than `full_scan_threshold` additional indexing won't be used - in this case full-scan search should be preferred by query planner and additional indexing is not required.",
"type": "integer",
"format": "uint",
"minimum": 0
},
"m": {
"description": "Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.",
"type": "integer",
"format": "uint",
"minimum": 0
}
]
}
},
"HnswConfigDiff": {
"type": "object",
"properties": {
"ef_construct": {
"description": "Number of neighbours to consider during the index building. Larger the value - more accurate the search, more time required to build index.",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
},
"full_scan_threshold": {
"description": "Minimal amount of points for additional payload-based indexing. If payload chunk is smaller than `full_scan_threshold` additional indexing won't be used - in this case full-scan search should be preferred by query planner and additional indexing is not required.",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
},
"m": {
"description": "Number of edges per node in the index graph. Larger the value - more accurate the search, more space required.",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
}
}
},
"Match": {
"type": "object",
@@ -427,42 +480,131 @@
}
}
},
"PayloadIndexType": {
"description": "Type of payload index",
"OptimizersConfig": {
"type": "object",
"required": [
"deleted_threshold",
"flush_interval_sec",
"indexing_threshold",
"max_segment_number",
"memmap_threshold",
"payload_indexing_threshold",
"vacuum_min_vector_number"
],
"properties": {
"deleted_threshold": {
"description": "The minimal fraction of deleted vectors in a segment, required to perform segment optimization",
"type": "number",
"format": "double"
},
"flush_interval_sec": {
"description": "Minimum interval between forced flushes.",
"type": "integer",
"format": "uint64",
"minimum": 0
},
"indexing_threshold": {
"description": "Maximum number of vectors allowed for plain index. Default value based on https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md",
"type": "integer",
"format": "uint",
"minimum": 0
},
"max_segment_number": {
"description": "If the number of segments exceeds this value, the optimizer will merge the smallest segments.",
"type": "integer",
"format": "uint",
"minimum": 0
},
"memmap_threshold": {
"description": "Maximum number of vectors to store in-memory per segment. Segments larger than this threshold will be stored as read-only memmaped file.",
"type": "integer",
"format": "uint",
"minimum": 0
},
"payload_indexing_threshold": {
"description": "Starting from this amount of vectors per-segment the engine will start building index for payload.",
"type": "integer",
"format": "uint",
"minimum": 0
},
"vacuum_min_vector_number": {
"description": "The minimal number of vectors in a segment, required to perform segment optimization",
"type": "integer",
"format": "uint",
"minimum": 0
}
}
},
"OptimizersConfigDiff": {
"type": "object",
"properties": {
"deleted_threshold": {
"description": "The minimal fraction of deleted vectors in a segment, required to perform segment optimization",
"type": "number",
"format": "double",
"nullable": true
},
"flush_interval_sec": {
"description": "Minimum interval between forced flushes.",
"type": "integer",
"format": "uint64",
"minimum": 0,
"nullable": true
},
"indexing_threshold": {
"description": "Maximum number of vectors allowed for plain index. Default value based on https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
},
"max_segment_number": {
"description": "If the number of segments exceeds this value, the optimizer will merge the smallest segments.",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
},
"memmap_threshold": {
"description": "Maximum number of vectors to store in-memory per segment. Segments larger than this threshold will be stored as read-only memmaped file.",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
},
"payload_indexing_threshold": {
"description": "Starting from this amount of vectors per-segment the engine will start building index for payload.",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
},
"vacuum_min_vector_number": {
"description": "The minimal number of vectors in a segment, required to perform segment optimization",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
}
}
},
"PayloadInterface": {
"anyOf": [
{
"description": "Do not index anything, just keep of what should be indexed later",
"type": "object",
"required": [
"type"
],
"properties": {
"type": {
"type": "string",
"enum": [
"plain"
]
}
}
"$ref": "#/components/schemas/PayloadVariant_for_String"
},
{
"description": "Build payload index. Index is saved on disc, but index itself is in RAM",
"type": "object",
"required": [
"type"
],
"properties": {
"type": {
"type": "string",
"enum": [
"struct"
]
}
}
"$ref": "#/components/schemas/PayloadVariant_for_int64"
},
{
"$ref": "#/components/schemas/PayloadVariant_for_double"
},
{
"$ref": "#/components/schemas/PayloadInterfaceStrict"
}
]
},
"PayloadInterface": {
"PayloadInterfaceStrict": {
"anyOf": [
{
"type": "object",
@@ -572,7 +714,8 @@
}
}
}
}
},
"additionalProperties": false
},
{
"description": "Deletes specified payload values if they are assigned",
@@ -605,7 +748,8 @@
}
}
}
}
},
"additionalProperties": false
},
{
"description": "Drops all Payload values associated with given points.",
@@ -630,7 +774,8 @@
}
}
}
}
},
"additionalProperties": false
}
]
},
@@ -827,7 +972,8 @@
}
}
}
}
},
"additionalProperties": false
},
{
"description": "Insert points from a list",
@@ -842,7 +988,8 @@
"$ref": "#/components/schemas/PointStruct"
}
}
}
},
"additionalProperties": false
}
]
},
@@ -858,7 +1005,8 @@
"upsert_points": {
"$ref": "#/components/schemas/PointInsertOperations"
}
}
},
"additionalProperties": false
},
{
"description": "Delete point if exists",
@@ -883,7 +1031,8 @@
}
}
}
}
},
"additionalProperties": false
}
]
},
@@ -1074,31 +1223,16 @@
},
"SearchParams": {
"description": "Additional parameters of the search",
"anyOf": [
{
"description": "Params relevant to HNSW index",
"type": "object",
"required": [
"hnsw"
],
"properties": {
"hnsw": {
"type": "object",
"required": [
"ef"
],
"properties": {
"ef": {
"description": "Size of the beam in a beam-search. Larger the value - more accurate the result, more time required for search.",
"type": "integer",
"format": "uint",
"minimum": 0
}
}
}
}
"type": "object",
"properties": {
"hnsw_ef": {
"description": "Params relevant to HNSW index /// Size of the beam in a beam-search. Larger the value - more accurate the result, more time required for search.",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
}
]
}
},
"SearchRequest": {
"description": "Search request",
@@ -1146,43 +1280,6 @@
}
}
},
"SegmentConfig": {
"type": "object",
"required": [
"distance",
"index",
"storage_type",
"vector_size"
],
"properties": {
"distance": {
"$ref": "#/components/schemas/Distance"
},
"index": {
"$ref": "#/components/schemas/Indexes"
},
"payload_index": {
"description": "Payload Indexes",
"anyOf": [
{
"$ref": "#/components/schemas/PayloadIndexType"
},
{
"nullable": true
}
]
},
"storage_type": {
"$ref": "#/components/schemas/StorageType"
},
"vector_size": {
"description": "Size of a vectors used",
"type": "integer",
"format": "uint",
"minimum": 0
}
}
},
"StorageOperations": {
"anyOf": [
{
@@ -1203,10 +1300,11 @@
"distance": {
"$ref": "#/components/schemas/Distance"
},
"index": {
"hnsw_config": {
"description": "Custom params for HNSW index. If none - values from service configuration file are used.",
"anyOf": [
{
"$ref": "#/components/schemas/Indexes"
"$ref": "#/components/schemas/HnswConfigDiff"
},
{
"nullable": true
@@ -1216,14 +1314,69 @@
"name": {
"type": "string"
},
"optimizers_config": {
"description": "Custom params for Optimizers. If none - values from service configuration file are used.",
"anyOf": [
{
"$ref": "#/components/schemas/OptimizersConfigDiff"
},
{
"nullable": true
}
]
},
"vector_size": {
"type": "integer",
"format": "uint",
"minimum": 0
},
"wal_config": {
"description": "Custom params for WAL. If none - values from service configuration file are used.",
"anyOf": [
{
"$ref": "#/components/schemas/WalConfigDiff"
},
{
"nullable": true
}
]
}
}
}
}
},
"additionalProperties": false
},
{
"description": "Update parameters of the existing collection",
"type": "object",
"required": [
"update_collection"
],
"properties": {
"update_collection": {
"type": "object",
"required": [
"name"
],
"properties": {
"name": {
"type": "string"
},
"optimizers_config": {
"description": "Custom params for Optimizers. If none - values from service configuration file are used. This operation is blocking, it will only proceed ones all current optimizations are complete",
"anyOf": [
{
"$ref": "#/components/schemas/OptimizersConfigDiff"
},
{
"nullable": true
}
]
}
}
}
},
"additionalProperties": false
},
{
"description": "Delete collection with given name",
@@ -1235,7 +1388,8 @@
"delete_collection": {
"type": "string"
}
}
},
"additionalProperties": false
},
{
"description": "Perform changes of collection aliases. Alias changes are atomic, meaning that no collection modifications can happen between alias operations.",
@@ -1258,42 +1412,8 @@
}
}
}
}
}
]
},
"StorageType": {
"description": "Type of vector storage",
"anyOf": [
{
"description": "Store vectors in memory and use persistence storage only if vectors are changed",
"type": "object",
"required": [
"type"
],
"properties": {
"type": {
"type": "string",
"enum": [
"in_memory"
]
}
}
},
{
"description": "Use memmap to store vectors, a little slower than `InMemory`, but requires little RAM",
"type": "object",
"required": [
"type"
],
"properties": {
"type": {
"type": "string",
"enum": [
"mmap"
]
}
}
},
"additionalProperties": false
}
]
},
@@ -1321,6 +1441,46 @@
"acknowledged",
"completed"
]
},
"WalConfig": {
"type": "object",
"required": [
"wal_capacity_mb",
"wal_segments_ahead"
],
"properties": {
"wal_capacity_mb": {
"description": "Size of a single WAL segment in MB",
"type": "integer",
"format": "uint",
"minimum": 0
},
"wal_segments_ahead": {
"description": "Number of WAL segments to create ahead of actually used ones",
"type": "integer",
"format": "uint",
"minimum": 0
}
}
},
"WalConfigDiff": {
"type": "object",
"properties": {
"wal_capacity_mb": {
"description": "Size of a single WAL segment in MB",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
},
"wal_segments_ahead": {
"description": "Number of WAL segments to create ahead of actually used ones",
"type": "integer",
"format": "uint",
"minimum": 0,
"nullable": true
}
}
}
}
}

486
openapi/openapi-merged.yaml
View File

@@ -60,19 +60,33 @@ info:
{
"result": {
"status": "green",
"vectors_count": 0,
"segments_count": 5,
"disk_data_size": 0,
"ram_data_size": 0,
"config": {
"vector_size": 4,
"index": {
"type": "plain",
"options": {}
"params": {
"vector_size": 4,
"distance": "Dot"
},
"distance": "Dot",
"storage_type": {
"type": "in_memory"
"hnsw_config": {
"m": 16,
"ef_construct": 100,
"full_scan_threshold": 10000
},
"optimizer_config": {
"deleted_threshold": 0.2,
"vacuum_min_vector_number": 1000,
"max_segment_number": 5,
"memmap_threshold": 50000,
"indexing_threshold": 20000,
"payload_indexing_threshold": 10000,
"flush_interval_sec": 1
},
"wal_config": {
"wal_capacity_mb": 32,
"wal_segments_ahead": 0
}
}
},
@@ -563,7 +577,8 @@ components:
AliasOperations:
anyOf:
- description: Create alternative name for a collection. Collection will be
- additionalProperties: false
description: Create alternative name for a collection. Collection will be
available under both names for search, retrieve,
properties:
create_alias:
@@ -579,7 +594,8 @@ components:
required:
- create_alias
type: object
- description: Delete alias if exists
- additionalProperties: false
description: Delete alias if exists
properties:
delete_alias:
properties:
@@ -591,7 +607,8 @@ components:
required:
- delete_alias
type: object
- description: Change alias to a new one
- additionalProperties: false
description: Change alias to a new one
properties:
rename_alias:
properties:
@@ -606,6 +623,22 @@ components:
required:
- rename_alias
type: object
CollectionConfig:
properties:
hnsw_config:
$ref: '#/components/schemas/HnswConfig'
optimizer_config:
$ref: '#/components/schemas/OptimizersConfig'
params:
$ref: '#/components/schemas/CollectionParams'
wal_config:
$ref: '#/components/schemas/WalConfig'
required:
- hnsw_config
- optimizer_config
- params
- wal_config
type: object
CollectionDescription:
properties:
name:
@@ -617,7 +650,7 @@ components:
description: Current statistics and configuration of the collection.
properties:
config:
$ref: '#/components/schemas/SegmentConfig'
$ref: '#/components/schemas/CollectionConfig'
disk_data_size:
description: Disk space, used by collection
format: uint
@@ -633,6 +666,8 @@ components:
format: uint
minimum: 0
type: integer
status:
$ref: '#/components/schemas/CollectionStatus'
vectors_count:
description: Number of vectors in collection
format: uint
@@ -643,8 +678,28 @@ components:
- disk_data_size
- ram_data_size
- segments_count
- status
- vectors_count
type: object
CollectionParams:
properties:
distance:
$ref: '#/components/schemas/Distance'
vector_size:
description: Size of a vectors used
format: uint
minimum: 0
type: integer
required:
- distance
- vector_size
type: object
CollectionStatus:
enum:
- green
- yellow
- red
type: string
CollectionUpdateOperations:
anyOf:
- $ref: '#/components/schemas/PointOperations'
@@ -701,14 +756,16 @@ components:
type: object
FieldIndexOperations:
anyOf:
- description: Create index for payload field
- additionalProperties: false
description: Create index for payload field
properties:
create_index:
type: string
required:
- create_index
type: object
- description: Delete index for the field
- additionalProperties: false
description: Delete index for the field
properties:
delete_index:
type: string
@@ -779,55 +836,65 @@ components:
minimum: 0
type: integer
type: array
uniqueItems: true
required:
- has_id
type: object
Indexes:
anyOf:
- description: Do not use any index, scan whole vector collection during search.
Guarantee 100% precision, but may be time consuming on large collections.
properties:
options:
type: object
type:
enum:
- plain
type: string
required:
- options
- type
type: object
- description: Use filterable HNSW index for approximate search. Is very fast
even on a very huge collections, but require additional space to store index
and additional time to build it.
properties:
options:
properties:
ef_construct:
description: Number of neighbours to consider during the index building.
Larger the value - more accurate the search, more time required
to build index.
format: uint
minimum: 0
type: integer
m:
description: Number of edges per node in the index graph. Larger the
value - more accurate the search, more space required.
format: uint
minimum: 0
type: integer
required:
- ef_construct
- m
type: object
type:
enum:
- hnsw
type: string
required:
- options
- type
type: object
HnswConfig:
properties:
ef_construct:
description: Number of neighbours to consider during the index building.
Larger the value - more accurate the search, more time required to build
index.
format: uint
minimum: 0
type: integer
full_scan_threshold:
description: Minimal amount of points for additional payload-based indexing.
If payload chunk is smaller than `full_scan_threshold` additional indexing
won't be used - in this case full-scan search should be preferred by query
planner and additional indexing is not required.
format: uint
minimum: 0
type: integer
m:
description: Number of edges per node in the index graph. Larger the value
- more accurate the search, more space required.
format: uint
minimum: 0
type: integer
required:
- ef_construct
- full_scan_threshold
- m
type: object
HnswConfigDiff:
properties:
ef_construct:
description: Number of neighbours to consider during the index building.
Larger the value - more accurate the search, more time required to build
index.
format: uint
minimum: 0
nullable: true
type: integer
full_scan_threshold:
description: Minimal amount of points for additional payload-based indexing.
If payload chunk is smaller than `full_scan_threshold` additional indexing
won't be used - in this case full-scan search should be preferred by query
planner and additional indexing is not required.
format: uint
minimum: 0
nullable: true
type: integer
m:
description: Number of edges per node in the index graph. Larger the value
- more accurate the search, more space required.
format: uint
minimum: 0
nullable: true
type: integer
type: object
Match:
properties:
integer:
@@ -840,29 +907,114 @@ components:
nullable: true
type: string
type: object
PayloadIndexType:
anyOf:
- description: Do not index anything, just keep of what should be indexed later
properties:
type:
enum:
- plain
type: string
required:
- type
type: object
- description: Build payload index. Index is saved on disc, but index itself
is in RAM
properties:
type:
enum:
- struct
type: string
required:
- type
type: object
description: Type of payload index
OptimizersConfig:
properties:
deleted_threshold:
description: The minimal fraction of deleted vectors in a segment, required
to perform segment optimization
format: double
type: number
flush_interval_sec:
description: Minimum interval between forced flushes.
format: uint64
minimum: 0
type: integer
indexing_threshold:
description: Maximum number of vectors allowed for plain index. Default
value based on https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md
format: uint
minimum: 0
type: integer
max_segment_number:
description: If the number of segments exceeds this value, the optimizer
will merge the smallest segments.
format: uint
minimum: 0
type: integer
memmap_threshold:
description: Maximum number of vectors to store in-memory per segment. Segments
larger than this threshold will be stored as read-only memmaped file.
format: uint
minimum: 0
type: integer
payload_indexing_threshold:
description: Starting from this amount of vectors per-segment the engine
will start building index for payload.
format: uint
minimum: 0
type: integer
vacuum_min_vector_number:
description: The minimal number of vectors in a segment, required to perform
segment optimization
format: uint
minimum: 0
type: integer
required:
- deleted_threshold
- flush_interval_sec
- indexing_threshold
- max_segment_number
- memmap_threshold
- payload_indexing_threshold
- vacuum_min_vector_number
type: object
OptimizersConfigDiff:
properties:
deleted_threshold:
description: The minimal fraction of deleted vectors in a segment, required
to perform segment optimization
format: double
nullable: true
type: number
flush_interval_sec:
description: Minimum interval between forced flushes.
format: uint64
minimum: 0
nullable: true
type: integer
indexing_threshold:
description: Maximum number of vectors allowed for plain index. Default
value based on https://github.com/google-research/google-research/blob/master/scann/docs/algorithms.md
format: uint
minimum: 0
nullable: true
type: integer
max_segment_number:
description: If the number of segments exceeds this value, the optimizer
will merge the smallest segments.
format: uint
minimum: 0
nullable: true
type: integer
memmap_threshold:
description: Maximum number of vectors to store in-memory per segment. Segments
larger than this threshold will be stored as read-only memmaped file.
format: uint
minimum: 0
nullable: true
type: integer
payload_indexing_threshold:
description: Starting from this amount of vectors per-segment the engine
will start building index for payload.
format: uint
minimum: 0
nullable: true
type: integer
vacuum_min_vector_number:
description: The minimal number of vectors in a segment, required to perform
segment optimization
format: uint
minimum: 0
nullable: true
type: integer
type: object
PayloadInterface:
anyOf:
- $ref: '#/components/schemas/PayloadVariant_for_String'
- $ref: '#/components/schemas/PayloadVariant_for_int64'
- $ref: '#/components/schemas/PayloadVariant_for_double'
- $ref: '#/components/schemas/PayloadInterfaceStrict'
PayloadInterfaceStrict:
anyOf:
- properties:
type:
@@ -910,7 +1062,8 @@ components:
type: object
PayloadOps:
anyOf:
- description: Set payload value, overrides if it is already exists
- additionalProperties: false
description: Set payload value, overrides if it is already exists
properties:
set_payload:
properties:
@@ -932,7 +1085,8 @@ components:
required:
- set_payload
type: object
- description: Deletes specified payload values if they are assigned
- additionalProperties: false
description: Deletes specified payload values if they are assigned
properties:
delete_payload:
properties:
@@ -954,7 +1108,8 @@ components:
required:
- delete_payload
type: object
- description: Drops all Payload values associated with given points.
- additionalProperties: false
description: Drops all Payload values associated with given points.
properties:
clear_payload:
properties:
@@ -1057,7 +1212,8 @@ components:
type: array
PointInsertOperations:
anyOf:
- description: Inset points from a batch.
- additionalProperties: false
description: Inset points from a batch.
properties:
batch:
properties:
@@ -1089,7 +1245,8 @@ components:
required:
- batch
type: object
- description: Insert points from a list
- additionalProperties: false
description: Insert points from a list
properties:
points:
items:
@@ -1100,14 +1257,16 @@ components:
type: object
PointOperations:
anyOf:
- description: Insert or update points
- additionalProperties: false
description: Insert or update points
properties:
upsert_points:
$ref: '#/components/schemas/PointInsertOperations'
required:
- upsert_points
type: object
- description: Delete point if exists
- additionalProperties: false
description: Delete point if exists
properties:
delete_points:
properties:
@@ -1257,24 +1416,16 @@ components:
- score
type: object
SearchParams:
anyOf:
- description: Params relevant to HNSW index
properties:
hnsw:
properties:
ef:
description: Size of the beam in a beam-search. Larger the value -
more accurate the result, more time required for search.
format: uint
minimum: 0
type: integer
required:
- ef
type: object
required:
- hnsw
type: object
description: Additional parameters of the search
properties:
hnsw_ef:
description: Params relevant to HNSW index /// Size of the beam in a beam-search.
Larger the value - more accurate the result, more time required for search.
format: uint
minimum: 0
nullable: true
type: integer
type: object
SearchRequest:
description: Search request
properties:
@@ -1303,48 +1454,39 @@ components:
- top
- vector
type: object
SegmentConfig:
properties:
distance:
$ref: '#/components/schemas/Distance'
index:
$ref: '#/components/schemas/Indexes'
payload_index:
anyOf:
- $ref: '#/components/schemas/PayloadIndexType'
- nullable: true
description: Payload Indexes
storage_type:
$ref: '#/components/schemas/StorageType'
vector_size:
description: Size of a vectors used
format: uint
minimum: 0
type: integer
required:
- distance
- index
- storage_type
- vector_size
type: object
StorageOperations:
anyOf:
- description: Create new collection and (optionally) specify index params
- additionalProperties: false
description: Create new collection and (optionally) specify index params
properties:
create_collection:
properties:
distance:
$ref: '#/components/schemas/Distance'
index:
hnsw_config:
anyOf:
- $ref: '#/components/schemas/Indexes'
- $ref: '#/components/schemas/HnswConfigDiff'
- nullable: true
description: Custom params for HNSW index. If none - values from service
configuration file are used.
name:
type: string
optimizers_config:
anyOf:
- $ref: '#/components/schemas/OptimizersConfigDiff'
- nullable: true
description: Custom params for Optimizers. If none - values from
service configuration file are used.
vector_size:
format: uint
minimum: 0
type: integer
wal_config:
anyOf:
- $ref: '#/components/schemas/WalConfigDiff'
- nullable: true
description: Custom params for WAL. If none - values from service
configuration file are used.
required:
- distance
- name
@@ -1353,14 +1495,36 @@ components:
required:
- create_collection
type: object
- description: Delete collection with given name
- additionalProperties: false
description: Update parameters of the existing collection
properties:
update_collection:
properties:
name:
type: string
optimizers_config:
anyOf:
- $ref: '#/components/schemas/OptimizersConfigDiff'
- nullable: true
description: Custom params for Optimizers. If none - values from
service configuration file are used. This operation is blocking,
it will only proceed ones all current optimizations are complete
required:
- name
type: object
required:
- update_collection
type: object
- additionalProperties: false
description: Delete collection with given name
properties:
delete_collection:
type: string
required:
- delete_collection
type: object
- description: Perform changes of collection aliases. Alias changes are atomic,
- additionalProperties: false
description: Perform changes of collection aliases. Alias changes are atomic,
meaning that no collection modifications can happen between alias operations.
properties:
change_aliases:
@@ -1375,29 +1539,6 @@ components:
required:
- change_aliases
type: object
StorageType:
anyOf:
- description: Store vectors in memory and use persistence storage only if vectors
are changed
properties:
type:
enum:
- in_memory
type: string
required:
- type
type: object
- description: Use memmap to store vectors, a little slower than `InMemory`,
but requires little RAM
properties:
type:
enum:
- mmap
type: string
required:
- type
type: object
description: Type of vector storage
UpdateResult:
properties:
operation_id:
@@ -1416,3 +1557,34 @@ components:
- acknowledged
- completed
type: string
WalConfig:
properties:
wal_capacity_mb:
description: Size of a single WAL segment in MB
format: uint
minimum: 0
type: integer
wal_segments_ahead:
description: Number of WAL segments to create ahead of actually used ones
format: uint
minimum: 0
type: integer
required:
- wal_capacity_mb
- wal_segments_ahead
type: object
WalConfigDiff:
properties:
wal_capacity_mb:
description: Size of a single WAL segment in MB
format: uint
minimum: 0
nullable: true
type: integer
wal_segments_ahead:
description: Number of WAL segments to create ahead of actually used ones
format: uint
minimum: 0
nullable: true
type: integer
type: object

28
openapi/openapi.yaml
View File

@@ -60,19 +60,33 @@ info:
{
"result": {
"status": "green",
"vectors_count": 0,
"segments_count": 5,
"disk_data_size": 0,
"ram_data_size": 0,
"config": {
"vector_size": 4,
"index": {
"type": "plain",
"options": {}
"params": {
"vector_size": 4,
"distance": "Dot"
},
"distance": "Dot",
"storage_type": {
"type": "in_memory"
"hnsw_config": {
"m": 16,
"ef_construct": 100,
"full_scan_threshold": 10000
},
"optimizer_config": {
"deleted_threshold": 0.2,
"vacuum_min_vector_number": 1000,
"max_segment_number": 5,
"memmap_threshold": 50000,
"indexing_threshold": 20000,
"payload_indexing_threshold": 10000,
"flush_interval_sec": 1
},
"wal_config": {
"wal_capacity_mb": 32,
"wal_segments_ahead": 0
}
}
},