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

Update sparse vector benchmarks (#4163)

Browse Source
This commit is contained in:
xzfc
2024-05-07 20:34:37 +00:00
committed by timvisee
parent 0a58b16e80
commit 4372ee6915
19 changed files with 987 additions and 170 deletions
Download Patch File Download Diff File Expand all files Collapse all files

264
Cargo.lock generated
View File

@@ -9,13 +9,13 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "57a7559404a7f3573127aab53c08ce37a6c6a315c374a31070f3c91cd1b4a7fe"
dependencies = [
"bitflags 1.3.2",
"bytes",
"bytes 1.6.0",
"futures-core",
"futures-sink",
"log",
"memchr",
"pin-project-lite",
"tokio",
"pin-project-lite 0.2.12",
"tokio 1.37.0",
"tokio-util",
]
@@ -45,7 +45,7 @@ dependencies = [
"actix-utils",
"actix-web",
"bitflags 2.4.1",
"bytes",
"bytes 1.6.0",
"derive_more",
"futures-core",
"http-range",
@@ -53,7 +53,7 @@ dependencies = [
"mime",
"mime_guess",
"percent-encoding",
"pin-project-lite",
"pin-project-lite 0.2.12",
"v_htmlescape",
]
@@ -72,7 +72,7 @@ dependencies = [
"base64 0.21.0",
"bitflags 2.4.1",
"brotli",
"bytes",
"bytes 1.6.0",
"bytestring",
"derive_more",
"encoding_rs",
@@ -87,11 +87,11 @@ dependencies = [
"local-channel",
"mime",
"percent-encoding",
"pin-project-lite",
"pin-project-lite 0.2.12",
"rand 0.8.5",
"sha1",
"smallvec",
"tokio",
"tokio 1.37.0",
"tokio-util",
"tracing",
"zstd 0.13.0",
@@ -116,7 +116,7 @@ dependencies = [
"actix-multipart-derive",
"actix-utils",
"actix-web",
"bytes",
"bytes 1.6.0",
"derive_more",
"futures-core",
"futures-util",
@@ -129,7 +129,7 @@ dependencies = [
"serde_json",
"serde_plain",
"tempfile",
"tokio",
"tokio 1.37.0",
]
[[package]]
@@ -165,7 +165,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7ea16c295198e958ef31930a6ef37d0fb64e9ca3b6116e6b93a8bdae96ee1000"
dependencies = [
"futures-core",
"tokio",
"tokio 1.37.0",
]
[[package]]
@@ -182,7 +182,7 @@ dependencies = [
"mio",
"num_cpus",
"socket2 0.4.9",
"tokio",
"tokio 1.37.0",
"tracing",
]
@@ -194,7 +194,7 @@ checksum = "3b894941f818cfdc7ccc4b9e60fa7e53b5042a2e8567270f9147d5591893373a"
dependencies = [
"futures-core",
"paste",
"pin-project-lite",
"pin-project-lite 0.2.12",
]
[[package]]
@@ -208,9 +208,9 @@ dependencies = [
"actix-utils",
"futures-core",
"impl-more",
"pin-project-lite",
"pin-project-lite 0.2.12",
"rustls-pki-types",
"tokio",
"tokio 1.37.0",
"tokio-rustls 0.25.0",
"tokio-util",
"tracing",
@@ -223,7 +223,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "88a1dcdff1466e3c2488e1cb5c36a71822750ad43839937f85d2f4d9f8b705d8"
dependencies = [
"local-waker",
"pin-project-lite",
"pin-project-lite 0.2.12",
]
[[package]]
@@ -243,7 +243,7 @@ dependencies = [
"actix-utils",
"actix-web-codegen",
"ahash",
"bytes",
"bytes 1.6.0",
"bytestring",
"cfg-if",
"cookie",
@@ -256,7 +256,7 @@ dependencies = [
"log",
"mime",
"once_cell",
"pin-project-lite",
"pin-project-lite 0.2.12",
"regex",
"serde",
"serde_json",
@@ -288,7 +288,7 @@ dependencies = [
"actix-web",
"futures-core",
"futures-util",
"pin-project-lite",
"pin-project-lite 0.2.12",
]
[[package]]
@@ -300,7 +300,7 @@ dependencies = [
"actix-http",
"actix-router",
"actix-web",
"bytes",
"bytes 1.6.0",
"futures",
"futures-util",
"log",
@@ -482,7 +482,7 @@ dependencies = [
"serde_json",
"sparse",
"thiserror",
"tokio",
"tokio 1.37.0",
"tonic",
"tonic-build",
"tracing",
@@ -520,6 +520,28 @@ version = "0.7.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8da52d66c7071e2e3fa2a1e5c6d088fec47b593032b254f5e980de8ea54454d6"
[[package]]
name = "async-compression"
version = "0.3.15"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "942c7cd7ae39e91bde4820d74132e9862e62c2f386c3aa90ccf55949f5bad63a"
dependencies = [
"brotli",
"bytes 0.5.6",
"bzip2",
"flate2",
"futures-core",
"futures-io",
"memchr",
"pin-project-lite 0.2.12",
"tokio 0.2.25",
"tokio 0.3.7",
"tokio 1.37.0",
"xz2",
"zstd 0.11.2+zstd.1.5.2",
"zstd-safe 5.0.2+zstd.1.5.2",
]
[[package]]
name = "async-stream"
version = "0.3.3"
@@ -584,7 +606,7 @@ dependencies = [
"async-trait",
"axum-core",
"bitflags 1.3.2",
"bytes",
"bytes 1.6.0",
"futures-util",
"http 0.2.9",
"http-body 0.4.5",
@@ -594,7 +616,7 @@ dependencies = [
"memchr",
"mime",
"percent-encoding",
"pin-project-lite",
"pin-project-lite 0.2.12",
"rustversion",
"serde",
"sync_wrapper",
@@ -610,7 +632,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b2f958c80c248b34b9a877a643811be8dbca03ca5ba827f2b63baf3a81e5fc4e"
dependencies = [
"async-trait",
"bytes",
"bytes 1.6.0",
"futures-util",
"http 0.2.9",
"http-body 0.4.5",
@@ -797,6 +819,12 @@ version = "1.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
[[package]]
name = "bytes"
version = "0.5.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0e4cec68f03f32e44924783795810fa50a7035d8c8ebe78580ad7e6c703fba38"
[[package]]
name = "bytes"
version = "1.6.0"
@@ -809,7 +837,7 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f7f83e57d9154148e355404702e2694463241880b939570d7c97c014da7a69a1"
dependencies = [
"bytes",
"bytes 1.6.0",
]
[[package]]
@@ -838,7 +866,7 @@ name = "cancel"
version = "0.0.0"
dependencies = [
"thiserror",
"tokio",
"tokio 1.37.0",
"tokio-util",
]
@@ -1055,7 +1083,7 @@ dependencies = [
"async-trait",
"atomicwrites",
"bitvec",
"bytes",
"bytes 1.6.0",
"cancel",
"chrono",
"collection",
@@ -1095,7 +1123,7 @@ dependencies = [
"tempfile",
"thiserror",
"tinyvec",
"tokio",
"tokio 1.37.0",
"tokio-util",
"tonic",
"tracing",
@@ -1132,7 +1160,7 @@ dependencies = [
"serde",
"thiserror",
"thread-priority",
"tokio",
"tokio 1.37.0",
"validator",
]
@@ -1199,7 +1227,7 @@ dependencies = [
"serde",
"serde_json",
"thread_local",
"tokio",
"tokio 1.37.0",
"tokio-stream",
"tonic",
"tracing",
@@ -1521,6 +1549,18 @@ dependencies = [
"parking_lot_core",
]
[[package]]
name = "dataset"
version = "0.0.0"
dependencies = [
"anyhow",
"async-compression",
"flate2",
"indicatif",
"project-root",
"reqwest",
]
[[package]]
name = "debugid"
version = "0.8.0"
@@ -1787,9 +1827,9 @@ checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
[[package]]
name = "flate2"
version = "1.0.28"
version = "1.0.30"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "46303f565772937ffe1d394a4fac6f411c6013172fadde9dcdb1e147a086940e"
checksum = "5f54427cfd1c7829e2a139fcefea601bf088ebca651d2bf53ebc600eac295dae"
dependencies = [
"crc32fast",
"miniz_oxide 0.7.1",
@@ -1967,7 +2007,7 @@ dependencies = [
"futures-sink",
"futures-task",
"memchr",
"pin-project-lite",
"pin-project-lite 0.2.12",
"pin-utils",
"slab",
]
@@ -2126,7 +2166,7 @@ version = "0.3.26"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "81fe527a889e1532da5c525686d96d4c2e74cdd345badf8dfef9f6b39dd5f5e8"
dependencies = [
"bytes",
"bytes 1.6.0",
"fnv",
"futures-core",
"futures-sink",
@@ -2134,7 +2174,7 @@ dependencies = [
"http 0.2.9",
"indexmap 2.2.6",
"slab",
"tokio",
"tokio 1.37.0",
"tokio-util",
"tracing",
]
@@ -2145,7 +2185,7 @@ version = "0.4.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "816ec7294445779408f36fe57bc5b7fc1cf59664059096c65f905c1c61f58069"
dependencies = [
"bytes",
"bytes 1.6.0",
"fnv",
"futures-core",
"futures-sink",
@@ -2153,7 +2193,7 @@ dependencies = [
"http 1.0.0",
"indexmap 2.2.6",
"slab",
"tokio",
"tokio 1.37.0",
"tokio-util",
"tracing",
]
@@ -2266,7 +2306,7 @@ version = "0.2.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bd6effc99afb63425aff9b05836f029929e345a6148a14b7ecd5ab67af944482"
dependencies = [
"bytes",
"bytes 1.6.0",
"fnv",
"itoa",
]
@@ -2277,7 +2317,7 @@ version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b32afd38673a8016f7c9ae69e5af41a58f81b1d31689040f2f1959594ce194ea"
dependencies = [
"bytes",
"bytes 1.6.0",
"fnv",
"itoa",
]
@@ -2288,9 +2328,9 @@ version = "0.4.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d5f38f16d184e36f2408a55281cd658ecbd3ca05cce6d6510a176eca393e26d1"
dependencies = [
"bytes",
"bytes 1.6.0",
"http 0.2.9",
"pin-project-lite",
"pin-project-lite 0.2.12",
]
[[package]]
@@ -2299,7 +2339,7 @@ version = "1.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1cac85db508abc24a2e48553ba12a996e87244a0395ce011e62b37158745d643"
dependencies = [
"bytes",
"bytes 1.6.0",
"http 1.0.0",
]
@@ -2309,11 +2349,11 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0475f8b2ac86659c21b64320d5d653f9efe42acd2a4e560073ec61a155a34f1d"
dependencies = [
"bytes",
"bytes 1.6.0",
"futures-core",
"http 1.0.0",
"http-body 1.0.0",
"pin-project-lite",
"pin-project-lite 0.2.12",
]
[[package]]
@@ -2356,7 +2396,7 @@ version = "0.14.26"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ab302d72a6f11a3b910431ff93aae7e773078c769f0a3ef15fb9ec692ed147d4"
dependencies = [
"bytes",
"bytes 1.6.0",
"futures-channel",
"futures-core",
"futures-util",
@@ -2366,9 +2406,9 @@ dependencies = [
"httparse",
"httpdate",
"itoa",
"pin-project-lite",
"pin-project-lite 0.2.12",
"socket2 0.4.9",
"tokio",
"tokio 1.37.0",
"tower-service",
"tracing",
"want",
@@ -2380,7 +2420,7 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "186548d73ac615b32a73aafe38fb4f56c0d340e110e5a200bcadbaf2e199263a"
dependencies = [
"bytes",
"bytes 1.6.0",
"futures-channel",
"futures-util",
"h2 0.4.4",
@@ -2388,9 +2428,9 @@ dependencies = [
"http-body 1.0.0",
"httparse",
"itoa",
"pin-project-lite",
"pin-project-lite 0.2.12",
"smallvec",
"tokio",
"tokio 1.37.0",
"want",
]
@@ -2406,7 +2446,7 @@ dependencies = [
"hyper-util",
"rustls 0.22.4",
"rustls-pki-types",
"tokio",
"tokio 1.37.0",
"tokio-rustls 0.25.0",
"tower-service",
]
@@ -2418,8 +2458,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bbb958482e8c7be4bc3cf272a766a2b0bf1a6755e7a6ae777f017a31d11b13b1"
dependencies = [
"hyper 0.14.26",
"pin-project-lite",
"tokio",
"pin-project-lite 0.2.12",
"tokio 1.37.0",
"tokio-io-timeout",
]
@@ -2429,15 +2469,15 @@ version = "0.1.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ca38ef113da30126bbff9cd1705f9273e15d45498615d138b0c20279ac7a76aa"
dependencies = [
"bytes",
"bytes 1.6.0",
"futures-channel",
"futures-util",
"http 1.0.0",
"http-body 1.0.0",
"hyper 1.2.0",
"pin-project-lite",
"pin-project-lite 0.2.12",
"socket2 0.5.5",
"tokio",
"tokio 1.37.0",
"tower",
"tower-service",
"tracing",
@@ -3101,6 +3141,17 @@ dependencies = [
"tracing-subscriber",
]
[[package]]
name = "lzma-sys"
version = "0.1.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5fda04ab3764e6cde78b9974eec4f779acaba7c4e84b36eca3cf77c581b85d27"
dependencies = [
"cc",
"libc",
"pkg-config",
]
[[package]]
name = "macro_rules_attribute"
version = "0.2.0"
@@ -3653,6 +3704,12 @@ dependencies = [
"syn 1.0.107",
]
[[package]]
name = "pin-project-lite"
version = "0.1.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "257b64915a082f7811703966789728173279bdebb956b143dbcd23f6f970a777"
[[package]]
name = "pin-project-lite"
version = "0.2.12"
@@ -3818,6 +3875,12 @@ dependencies = [
"hex",
]
[[package]]
name = "project-root"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8bccbff07d5ed689c4087d20d7307a52ab6141edeedf487c3876a55b86cf63df"
[[package]]
name = "prometheus"
version = "0.13.4"
@@ -3858,7 +3921,7 @@ version = "0.11.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b82eaa1d779e9a4bc1c3217db8ffbeabaae1dca241bf70183242128d48681cd"
dependencies = [
"bytes",
"bytes 1.6.0",
"prost-derive 0.11.9",
]
@@ -3868,7 +3931,7 @@ version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "146c289cda302b98a28d40c8b3b90498d6e526dd24ac2ecea73e4e491685b94a"
dependencies = [
"bytes",
"bytes 1.6.0",
"prost-derive 0.12.3",
]
@@ -3878,7 +3941,7 @@ version = "0.11.9"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "119533552c9a7ffacc21e099c24a0ac8bb19c2a2a3f363de84cd9b844feab270"
dependencies = [
"bytes",
"bytes 1.6.0",
"heck 0.4.1",
"itertools 0.10.5",
"lazy_static",
@@ -3902,7 +3965,7 @@ version = "0.12.3"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c55e02e35260070b6f716a2423c2ff1c3bb1642ddca6f99e1f26d06268a0e2d2"
dependencies = [
"bytes",
"bytes 1.6.0",
"heck 0.4.1",
"itertools 0.10.5",
"log",
@@ -4101,7 +4164,7 @@ dependencies = [
"tempfile",
"thiserror",
"tikv-jemallocator",
"tokio",
"tokio 1.37.0",
"tonic",
"tonic-reflection",
"tower",
@@ -4364,7 +4427,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "566cafdd92868e0939d3fb961bd0dc25fcfaaed179291093b3d43e6b3150ea10"
dependencies = [
"base64 0.22.0",
"bytes",
"bytes 1.6.0",
"futures-channel",
"futures-core",
"futures-util",
@@ -4381,7 +4444,7 @@ dependencies = [
"mime",
"once_cell",
"percent-encoding",
"pin-project-lite",
"pin-project-lite 0.2.12",
"rustls 0.22.4",
"rustls-pemfile 2.1.2",
"rustls-pki-types",
@@ -4389,7 +4452,7 @@ dependencies = [
"serde_json",
"serde_urlencoded",
"sync_wrapper",
"tokio",
"tokio 1.37.0",
"tokio-rustls 0.25.0",
"tokio-util",
"tower-service",
@@ -4865,12 +4928,14 @@ dependencies = [
"chrono",
"common",
"criterion",
"dataset",
"fnv",
"fs_extra",
"generic-tests",
"geo",
"geohash",
"indexmap 2.2.6",
"indicatif",
"io",
"io-uring",
"is_sorted",
@@ -5213,15 +5278,21 @@ name = "sparse"
version = "0.1.0"
dependencies = [
"common",
"criterion",
"dataset",
"indicatif",
"io",
"itertools 0.12.1",
"memmap2 0.9.4",
"memory",
"ordered-float 4.2.0",
"parking_lot",
"pprof",
"rand 0.8.5",
"schemars",
"serde",
"serde_json",
"sparse",
"tempfile",
"validator",
]
@@ -5286,7 +5357,7 @@ dependencies = [
"tar",
"tempfile",
"thiserror",
"tokio",
"tokio 1.37.0",
"tonic",
"tracing",
"url",
@@ -5602,6 +5673,26 @@ version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cda74da7e1a664f795bb1f8a87ec406fb89a02522cf6e50620d016add6dbbf5c"
[[package]]
name = "tokio"
version = "0.2.25"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6703a273949a90131b290be1fe7b039d0fc884aa1935860dfcbe056f28cd8092"
dependencies = [
"bytes 0.5.6",
"pin-project-lite 0.1.12",
]
[[package]]
name = "tokio"
version = "0.3.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "46409491c9375a693ce7032101970a54f8a2010efb77e13f70788f0d84489e39"
dependencies = [
"autocfg",
"pin-project-lite 0.2.12",
]
[[package]]
name = "tokio"
version = "1.37.0"
@@ -5609,12 +5700,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1adbebffeca75fcfd058afa480fb6c0b81e165a0323f9c9d39c9697e37c46787"
dependencies = [
"backtrace",
"bytes",
"bytes 1.6.0",
"libc",
"mio",
"num_cpus",
"parking_lot",
"pin-project-lite",
"pin-project-lite 0.2.12",
"signal-hook-registry",
"socket2 0.5.5",
"tokio-macros",
@@ -5628,8 +5719,8 @@ version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "30b74022ada614a1b4834de765f9bb43877f910cc8ce4be40e89042c9223a8bf"
dependencies = [
"pin-project-lite",
"tokio",
"pin-project-lite 0.2.12",
"tokio 1.37.0",
]
[[package]]
@@ -5650,7 +5741,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081"
dependencies = [
"rustls 0.21.11",
"tokio",
"tokio 1.37.0",
]
[[package]]
@@ -5661,7 +5752,7 @@ checksum = "775e0c0f0adb3a2f22a00c4745d728b479985fc15ee7ca6a2608388c5569860f"
dependencies = [
"rustls 0.22.4",
"rustls-pki-types",
"tokio",
"tokio 1.37.0",
]
[[package]]
@@ -5671,8 +5762,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d660770404473ccd7bc9f8b28494a811bc18542b915c0855c51e8f419d5223ce"
dependencies = [
"futures-core",
"pin-project-lite",
"tokio",
"pin-project-lite 0.2.12",
"tokio 1.37.0",
]
[[package]]
@@ -5681,11 +5772,11 @@ version = "0.7.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9cf6b47b3771c49ac75ad09a6162f53ad4b8088b76ac60e8ec1455b31a189fe1"
dependencies = [
"bytes",
"bytes 1.6.0",
"futures-core",
"futures-sink",
"pin-project-lite",
"tokio",
"pin-project-lite 0.2.12",
"tokio 1.37.0",
]
[[package]]
@@ -5740,7 +5831,7 @@ dependencies = [
"async-trait",
"axum",
"base64 0.21.0",
"bytes",
"bytes 1.6.0",
"flate2",
"futures-core",
"futures-util",
@@ -5753,7 +5844,7 @@ dependencies = [
"pin-project",
"prost 0.11.9",
"rustls-pemfile 1.0.3",
"tokio",
"tokio 1.37.0",
"tokio-rustls 0.24.1",
"tokio-stream",
"tower",
@@ -5783,7 +5874,7 @@ checksum = "0543d7092032041fbeac1f2c84304537553421a11a623c2301b12ef0264862c7"
dependencies = [
"prost 0.11.9",
"prost-types 0.11.9",
"tokio",
"tokio 1.37.0",
"tokio-stream",
"tonic",
]
@@ -5798,10 +5889,10 @@ dependencies = [
"futures-util",
"indexmap 1.9.2",
"pin-project",
"pin-project-lite",
"pin-project-lite 0.2.12",
"rand 0.8.5",
"slab",
"tokio",
"tokio 1.37.0",
"tokio-util",
"tower-layer",
"tower-service",
@@ -5827,7 +5918,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c3523ab5a71916ccf420eebdf5521fcef02141234bbc0b8a49f2fdc4544364ef"
dependencies = [
"log",
"pin-project-lite",
"pin-project-lite 0.2.12",
"tracing-attributes",
"tracing-core",
]
@@ -6635,6 +6726,15 @@ dependencies = [
"libc",
]
[[package]]
name = "xz2"
version = "0.1.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "388c44dc09d76f1536602ead6d325eb532f5c122f17782bd57fb47baeeb767e2"
dependencies = [
"lzma-sys",
]
[[package]]
name = "yada"
version = "0.5.0"

2
Cargo.toml
View File

@@ -123,7 +123,9 @@ fnv = "1.0"
futures = "0.3.30"
futures-util = "0.3.30"
indexmap = { version = "2", features = ["serde"] }
indicatif = "0.17.8"
parking_lot = { version = "0.12.2", features = ["deadlock_detection", "serde"] }
pprof = { version = "0.12", features = ["flamegraph", "prost-codec"] }
prost = "0.11.9"
prost-wkt-types = "0.4.2"
reqwest = { version = "0.12.4", default-features = false, features = ["http2", "stream", "rustls-tls", "blocking"] }

4
lib/collection/Cargo.toml
View File

@@ -21,7 +21,7 @@ collection = { path = ".", features = ["testing"] }
common = { path = "../common/common", features = ["testing"] }
[target.'cfg(not(target_os = "windows"))'.dev-dependencies]
pprof = { version = "0.12", features = ["flamegraph", "prost-codec"] }
pprof = { workspace = true }
[dependencies]
parking_lot = { workspace = true }
@@ -63,7 +63,7 @@ sparse = { path = "../sparse" }
api = { path = "../api" }
itertools = "0.12"
indicatif = "0.17.8"
indicatif = { workspace = true }
chrono = { workspace = true }
schemars = { workspace = true }
tar = { workspace = true }

18
lib/common/dataset/Cargo.toml Normal file
View File

@@ -0,0 +1,18 @@
[package]
name = "dataset"
version = "0.0.0"
authors = [
"Andrey Vasnetsov <andrey@qdrant.com>",
"Qdrant Team <info@qdrant.tech>",
]
license = "Apache-2.0"
edition = "2021"
publish = false
[dependencies]
anyhow = "1.0.82"
async-compression = {version="0.3.13", features=["all", "stream"]}
flate2 = { version = "1.0.30" }
indicatif = { workspace = true }
project-root = "0.2.2"
reqwest = { workspace = true }

108
lib/common/dataset/src/lib.rs Normal file
View File

@@ -0,0 +1,108 @@
use std::fs::{create_dir_all, File};
use std::path::{Path, PathBuf};
use anyhow::{anyhow, Context, Result};
use flate2::read::GzDecoder;
use indicatif::{ProgressBar, ProgressDrawTarget};
pub enum Dataset {
// https://github.com/qdrant/sparse-vectors-experiments
SpladeWikiMovies,
// https://github.com/qdrant/sparse-vectors-benchmark
NeurIps2023Full,
NeurIps2023_1M,
NeurIps2023Small,
NeurIps2023Queries,
}
impl Dataset {
pub fn download(&self) -> Result<PathBuf> {
download_cached(&self.url())
}
fn url(&self) -> String {
const NEUR_IPS_2023_BASE: &str =
"https://storage.googleapis.com/ann-challenge-sparse-vectors/csr";
match self {
Dataset::SpladeWikiMovies => {
"https://storage.googleapis.com/dataset-sparse-vectors/sparse-vectors.jsonl.gz"
.to_string()
}
Dataset::NeurIps2023Full => format!("{NEUR_IPS_2023_BASE}/base_full.csr.gz"),
Dataset::NeurIps2023_1M => format!("{NEUR_IPS_2023_BASE}/base_1M.csr.gz"),
Dataset::NeurIps2023Small => format!("{NEUR_IPS_2023_BASE}/base_small.csr.gz"),
Dataset::NeurIps2023Queries => format!("{NEUR_IPS_2023_BASE}/queries.dev.csr.gz"),
}
}
}
fn download_cached(url: &str) -> Result<PathBuf> {
// Filename without an ".gz" extension, e.g. "base_full.csr".
let basename = {
let path = Path::new(url);
match path.extension() {
Some(gz) if gz == "gz" => path.file_stem(),
_ => path.file_name(),
}
.ok_or_else(|| anyhow!("Failed to extract basename from {url}"))?
};
// Cache directory, e.g. "target/datasets".
let cache_dir = workspace_dir()
.join(std::env::var_os("CARGO_TARGET_DIR").unwrap_or_else(|| "target".into()))
.join("datasets");
// Cache file path, e.g. "target/datasets/base_full.csr".
let cache_path = cache_dir.join(basename);
if cache_path.exists() {
return Ok(cache_path);
}
eprintln!("Downloading {url} to {cache_path:?}...");
create_dir_all(cache_dir)?;
let resp = reqwest::blocking::get(url)?;
if !resp.status().is_success() {
anyhow::bail!("Failed to download {url}, status: {}", resp.status());
}
let total_size = resp.content_length();
// Download to a temporary file, e.g. "target/datasets/base_full.csr.tmp", to avoid
// incomplete files.
let mut tmp_fname = cache_path.clone().into_os_string();
tmp_fname.push(".tmp");
// Progress bar.
let pb = ProgressBar::with_draw_target(total_size, ProgressDrawTarget::stderr_with_hz(12));
pb.set_style(
indicatif::ProgressStyle::default_bar()
.template("{msg} {wide_bar} {bytes}/{total_bytes} (eta:{eta})")
.expect("failed to set style"),
);
// Download + decompress.
std::io::copy(
&mut GzDecoder::new(pb.wrap_read(resp)),
&mut File::create(&tmp_fname)?,
)?;
std::fs::rename(&tmp_fname, &cache_path)
.with_context(|| format!("Failed to rename {tmp_fname:?} to {cache_path:?}"))?;
Ok(cache_path)
}
fn workspace_dir() -> PathBuf {
let output = std::process::Command::new(env!("CARGO"))
.arg("locate-project")
.arg("--workspace")
.arg("--message-format=plain")
.output()
.unwrap()
.stdout;
let cargo_path = Path::new(std::str::from_utf8(&output).unwrap().trim());
cargo_path.parent().unwrap().to_path_buf()
}

4
lib/segment/Cargo.toml
View File

@@ -16,6 +16,8 @@ testing = []
[dev-dependencies]
criterion = "0.5"
dataset = { path = "../common/dataset" }
indicatif = { workspace = true }
rmp-serde = "~1.3"
rand_distr = "0.4.3"
walkdir = "2.5.0"
@@ -24,7 +26,7 @@ segment = { path = ".", features = ["testing"] }
proptest = "1.4.0"
[target.'cfg(not(target_os = "windows"))'.dev-dependencies]
pprof = { version = "0.12", features = ["flamegraph", "prost-codec"] }
pprof = { workspace = true }
[dependencies]
bitpacking = "0.9.2"

209
lib/segment/benches/sparse_index_search.rs
View File

@@ -6,10 +6,13 @@ use std::sync::Arc;
use common::cpu::CpuPermit;
use common::types::PointOffsetType;
use criterion::{criterion_group, criterion_main, Criterion};
use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use dataset::Dataset;
use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle};
use itertools::Itertools as _;
use rand::rngs::StdRng;
use rand::SeedableRng;
use segment::fixtures::sparse_fixtures::fixture_sparse_index_ram;
use segment::fixtures::sparse_fixtures::fixture_sparse_index_ram_from_iter;
use segment::index::hnsw_index::num_rayon_threads;
use segment::index::sparse_index::sparse_index_config::{SparseIndexConfig, SparseIndexType};
use segment::index::sparse_index::sparse_vector_index::SparseVectorIndex;
@@ -17,29 +20,67 @@ use segment::index::{PayloadIndex, VectorIndex};
use segment::types::PayloadSchemaType::Keyword;
use segment::types::{Condition, FieldCondition, Filter, Payload};
use serde_json::json;
use sparse::common::sparse_vector_fixture::random_positive_sparse_vector;
use sparse::common::sparse_vector::SparseVector;
use sparse::common::sparse_vector_fixture::{random_positive_sparse_vector, random_sparse_vector};
use sparse::index::inverted_index::inverted_index_mmap::InvertedIndexMmap;
use sparse::index::loaders::Csr;
use tempfile::Builder;
const NUM_VECTORS: usize = 50_000;
const MAX_SPARSE_DIM: usize = 30_000;
const NUM_QUERIES: usize = 2048;
const TOP: usize = 10;
const FULL_SCAN_THRESHOLD: usize = 1; // low value to trigger index usage by default
fn sparse_vector_index_search_benchmark(c: &mut Criterion) {
let mut group = c.benchmark_group("sparse-vector-search-group");
let mut rnd = StdRng::seed_from_u64(0);
let query_vectors = (0..NUM_QUERIES)
// Positive values to test pruning.
.map(|_| random_positive_sparse_vector(&mut rnd, MAX_SPARSE_DIM))
.collect::<Vec<_>>();
let mut rnd = StdRng::seed_from_u64(42);
let random_vectors = (0..NUM_VECTORS).map(|_| random_sparse_vector(&mut rnd, MAX_SPARSE_DIM));
sparse_vector_index_search_benchmark_impl(c, "random-50k", random_vectors, &query_vectors);
let dataset_vectors = Csr::open(Dataset::NeurIps2023_1M.download().unwrap()).unwrap();
let query_vectors = Csr::open(Dataset::NeurIps2023Queries.download().unwrap())
.unwrap()
.iter()
.map(|v| v.unwrap())
.collect_vec();
sparse_vector_index_search_benchmark_impl(
c,
"neurips2023-1M",
dataset_vectors.iter().map(|v| v.unwrap()),
&query_vectors,
);
}
fn sparse_vector_index_search_benchmark_impl(
c: &mut Criterion,
group: &str,
vectors: impl ExactSizeIterator<Item = SparseVector>,
query_vectors: &[SparseVector],
) {
let mut group = c.benchmark_group(format!("sparse_vector_index_search/{}", group));
group.sample_size(10);
let vectors_len = vectors.len();
let stopped = AtomicBool::new(false);
let mut rnd = StdRng::seed_from_u64(42);
let data_dir = Builder::new().prefix("data_dir").tempdir().unwrap();
let sparse_vector_index = fixture_sparse_index_ram(
&mut rnd,
NUM_VECTORS,
MAX_SPARSE_DIM,
let sparse_vector_index = fixture_sparse_index_ram_from_iter(
progress("Indexing (1/3)", vectors_len).wrap_iter(vectors),
FULL_SCAN_THRESHOLD,
data_dir.path(),
&stopped,
|| {
let pb = progress("Indexing (2/3)", vectors_len);
move || pb.inc(1)
},
);
// adding payload on field
@@ -59,11 +100,7 @@ fn sparse_vector_index_search_benchmark(c: &mut Criterion) {
}
drop(payload_index);
// shared query vector (positive values to test pruning)
let vector = random_positive_sparse_vector(&mut rnd, MAX_SPARSE_DIM);
eprintln!("sparse_vector size = {:#?}", vector.values.len());
let sparse_vector = vector.clone();
let query_vector = vector.into();
let mut query_vector_it = query_vectors.iter().cycle();
let permit_cpu_count = num_rayon_threads(0);
let permit = Arc::new(CpuPermit::dummy(permit_cpu_count as u32));
@@ -82,31 +119,41 @@ fn sparse_vector_index_search_benchmark(c: &mut Criterion) {
&stopped,
)
.unwrap();
let pb = progress("Indexing (3/3)", vectors_len);
sparse_vector_index_mmap
.build_index(permit, &stopped)
.build_index_with_progress(permit, &stopped, || pb.inc(1))
.unwrap();
assert_eq!(sparse_vector_index_mmap.indexed_vector_count(), NUM_VECTORS);
pb.finish_and_clear();
assert_eq!(sparse_vector_index_mmap.indexed_vector_count(), vectors_len);
// intent: bench `search` without filter on mmap inverted index
group.bench_function("mmap-inverted-index-search", |b| {
b.iter(|| {
let results = sparse_vector_index_mmap
.search(&[&query_vector], None, TOP, None, &Default::default())
.unwrap();
b.iter_batched(
|| query_vector_it.next().unwrap().clone().into(),
|vec| {
let results = sparse_vector_index_mmap
.search(&[&vec], None, TOP, None, &Default::default())
.unwrap();
assert_eq!(results[0].len(), TOP);
})
assert_eq!(results[0].len(), TOP);
},
BatchSize::SmallInput,
)
});
// intent: bench `search` without filter
group.bench_function("inverted-index-search", |b| {
b.iter(|| {
let results = sparse_vector_index
.search(&[&query_vector], None, TOP, None, &Default::default())
.unwrap();
b.iter_batched(
|| query_vector_it.next().unwrap().clone().into(),
|vec| {
let results = sparse_vector_index
.search(&[&vec], None, TOP, None, &Default::default())
.unwrap();
assert_eq!(results[0].len(), TOP);
})
assert_eq!(results[0].len(), TOP);
},
BatchSize::SmallInput,
)
});
// filter by field
@@ -116,22 +163,28 @@ fn sparse_vector_index_search_benchmark(c: &mut Criterion) {
)));
// intent: bench plain search when the filtered payload key is not indexed
group.bench_function("inverted-index-filtered-plain", |b| {
b.iter(|| {
let mut prefiltered_points = None;
let results = sparse_vector_index
.search_plain(
&sparse_vector,
&filter,
TOP,
&mut prefiltered_points,
&Default::default(),
)
.unwrap();
if vectors_len < 100_000 {
group.bench_function("inverted-index-filtered-plain", |b| {
b.iter_batched(
|| query_vector_it.next().unwrap(),
|vec| {
let mut prefiltered_points = None;
let results = sparse_vector_index
.search_plain(
vec,
&filter,
TOP,
&mut prefiltered_points,
&Default::default(),
)
.unwrap();
assert_eq!(results.len(), TOP);
})
});
assert_eq!(results.len(), TOP);
},
BatchSize::SmallInput,
)
});
}
let mut payload_index = sparse_vector_index.payload_index.borrow_mut();
@@ -144,42 +197,58 @@ fn sparse_vector_index_search_benchmark(c: &mut Criterion) {
// intent: bench `search` when the filtered payload key is indexed
group.bench_function("inverted-index-filtered-payload-index", |b| {
b.iter(|| {
let results = sparse_vector_index
.search(
&[&query_vector],
Some(&filter),
TOP,
None,
&Default::default(),
)
.unwrap();
b.iter_batched(
|| query_vector_it.next().unwrap().clone().into(),
|vec| {
let results = sparse_vector_index
.search(&[&vec], Some(&filter), TOP, None, &Default::default())
.unwrap();
assert_eq!(results[0].len(), TOP);
})
assert_eq!(results[0].len(), TOP);
},
BatchSize::SmallInput,
);
});
// intent: bench plain search when the filtered payload key is indexed
group.bench_function("plain-filtered-payload-index", |b| {
b.iter(|| {
let mut prefiltered_points = None;
let results = sparse_vector_index
.search_plain(
&sparse_vector,
&filter,
TOP,
&mut prefiltered_points,
&Default::default(),
)
.unwrap();
if vectors_len < 100_000 {
group.bench_function("plain-filtered-payload-index", |b| {
b.iter_batched(
|| query_vector_it.next().unwrap(),
|vec| {
let mut prefiltered_points = None;
let results = sparse_vector_index
.search_plain(
vec,
&filter,
TOP,
&mut prefiltered_points,
&Default::default(),
)
.unwrap();
assert_eq!(results.len(), TOP);
})
});
assert_eq!(results.len(), TOP);
},
BatchSize::SmallInput,
)
});
}
group.finish();
}
fn progress(name: &str, len: usize) -> ProgressBar {
let pb =
ProgressBar::with_draw_target(Some(len as u64), ProgressDrawTarget::stderr_with_hz(12));
pb.set_style(
ProgressStyle::default_bar()
.template("{msg} {wide_bar} {pos}/{len} (eta:{eta})")
.unwrap(),
);
pb.set_message(name.to_owned());
pb
}
#[cfg(not(target_os = "windows"))]
criterion_group! {
name = benches;

30
lib/segment/src/fixtures/sparse_fixtures.rs
View File

@@ -89,6 +89,24 @@ pub fn fixture_sparse_index_ram<R: Rng + ?Sized>(
data_dir: &Path,
stopped: &AtomicBool,
) -> SparseVectorIndex<InvertedIndexRam> {
fixture_sparse_index_ram_from_iter(
(0..num_vectors).map(|_| random_sparse_vector(rnd, max_dim)),
full_scan_threshold,
data_dir,
stopped,
|| || (),
)
}
/// Prepares a sparse vector index with a given iterator of sparse vectors
pub fn fixture_sparse_index_ram_from_iter<P: FnMut()>(
vectors: impl ExactSizeIterator<Item = SparseVector>,
full_scan_threshold: usize,
data_dir: &Path,
stopped: &AtomicBool,
progress: impl FnOnce() -> P,
) -> SparseVectorIndex<InvertedIndexRam> {
let num_vectors = vectors.len();
let mut sparse_vector_index = fixture_open_sparse_index(
data_dir,
num_vectors,
@@ -100,10 +118,9 @@ pub fn fixture_sparse_index_ram<R: Rng + ?Sized>(
let mut borrowed_storage = sparse_vector_index.vector_storage.borrow_mut();
// add points to storage
for idx in 0..num_vectors {
let vec = &random_sparse_vector(rnd, max_dim);
for (idx, vec) in vectors.enumerate() {
borrowed_storage
.insert_vector(idx as PointOffsetType, vec.into())
.insert_vector(idx as PointOffsetType, (&vec).into())
.unwrap();
}
drop(borrowed_storage);
@@ -114,7 +131,7 @@ pub fn fixture_sparse_index_ram<R: Rng + ?Sized>(
.vector_storage
.borrow()
.available_vector_count(),
num_vectors
num_vectors,
);
// assert no points are indexed following open for RAM index
@@ -124,7 +141,10 @@ pub fn fixture_sparse_index_ram<R: Rng + ?Sized>(
let permit = Arc::new(CpuPermit::dummy(permit_cpu_count as u32));
// build index to refresh RAM index
sparse_vector_index.build_index(permit, stopped).unwrap();
let tick_progress = progress();
sparse_vector_index
.build_index_with_progress(permit, stopped, tick_progress)
.unwrap();
assert_eq!(sparse_vector_index.indexed_vector_count(), num_vectors);
sparse_vector_index
}

7
lib/segment/src/index/hnsw_index/hnsw.rs
View File

@@ -665,7 +665,12 @@ impl<TGraphLinks: GraphLinks> VectorIndex for HNSWIndex<TGraphLinks> {
}
}
fn build_index(&mut self, permit: Arc<CpuPermit>, stopped: &AtomicBool) -> OperationResult<()> {
fn build_index_with_progress(
&mut self,
permit: Arc<CpuPermit>,
stopped: &AtomicBool,
_tick_progress: impl FnMut(),
) -> OperationResult<()> {
// Build main index graph
let id_tracker = self.id_tracker.borrow();
let vector_storage = self.vector_storage.borrow();

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

@@ -317,10 +317,11 @@ impl VectorIndex for PlainIndex {
}
}
fn build_index(
fn build_index_with_progress(
&mut self,
_permit: Arc<CpuPermit>,
_stopped: &AtomicBool,
_tick_progress: impl FnMut(),
) -> OperationResult<()> {
Ok(())
}

7
lib/segment/src/index/sparse_index/sparse_vector_index.rs
View File

@@ -71,6 +71,7 @@ impl<TInvertedIndex: InvertedIndex> SparseVectorIndex<TInvertedIndex> {
vector_storage.clone(),
path,
stopped,
|| (),
)?;
(config, inverted_index, indices_tracker)
} else if config_path.exists() {
@@ -114,6 +115,7 @@ impl<TInvertedIndex: InvertedIndex> SparseVectorIndex<TInvertedIndex> {
vector_storage: Arc<AtomicRefCell<VectorStorageEnum>>,
path: &Path,
stopped: &AtomicBool,
mut tick_progress: impl FnMut(),
) -> OperationResult<(TInvertedIndex, IndicesTracker)> {
let borrowed_vector_storage = vector_storage.borrow();
let borrowed_id_tracker = id_tracker.borrow();
@@ -143,6 +145,7 @@ impl<TInvertedIndex: InvertedIndex> SparseVectorIndex<TInvertedIndex> {
ram_index_builder.add(id, vector);
}
}
tick_progress();
}
Ok((
TInvertedIndex::from_ram_index(ram_index_builder.build(), path)?,
@@ -440,16 +443,18 @@ impl<TInvertedIndex: InvertedIndex> VectorIndex for SparseVectorIndex<TInvertedI
Ok(results)
}
fn build_index(
fn build_index_with_progress(
&mut self,
_permit: Arc<CpuPermit>,
stopped: &AtomicBool,
tick_progress: impl FnMut(),
) -> OperationResult<()> {
let (inverted_index, indices_tracker) = Self::build_inverted_index(
self.id_tracker.clone(),
self.vector_storage.clone(),
&self.path,
stopped,
tick_progress,
)?;
self.inverted_index = inverted_index;

39
lib/segment/src/index/vector_index_base.rs
View File

@@ -30,7 +30,17 @@ pub trait VectorIndex {
) -> OperationResult<Vec<Vec<ScoredPointOffset>>>;
/// Force internal index rebuild.
fn build_index(&mut self, permit: Arc<CpuPermit>, stopped: &AtomicBool) -> OperationResult<()>;
fn build_index(&mut self, permit: Arc<CpuPermit>, stopped: &AtomicBool) -> OperationResult<()> {
self.build_index_with_progress(permit, stopped, || ())
}
/// Force internal index rebuild.
fn build_index_with_progress(
&mut self,
permit: Arc<CpuPermit>,
stopped: &AtomicBool,
tick_progress: impl FnMut(),
) -> OperationResult<()>;
fn get_telemetry_data(&self, detail: TelemetryDetail) -> VectorIndexSearchesTelemetry;
@@ -91,13 +101,28 @@ impl VectorIndex for VectorIndexEnum {
}
}
fn build_index(&mut self, permit: Arc<CpuPermit>, stopped: &AtomicBool) -> OperationResult<()> {
fn build_index_with_progress(
&mut self,
permit: Arc<CpuPermit>,
stopped: &AtomicBool,
tick_progress: impl FnMut(),
) -> OperationResult<()> {
match self {
VectorIndexEnum::Plain(index) => index.build_index(permit, stopped),
VectorIndexEnum::HnswRam(index) => index.build_index(permit, stopped),
VectorIndexEnum::HnswMmap(index) => index.build_index(permit, stopped),
VectorIndexEnum::SparseRam(index) => index.build_index(permit, stopped),
VectorIndexEnum::SparseMmap(index) => index.build_index(permit, stopped),
VectorIndexEnum::Plain(index) => {
index.build_index_with_progress(permit, stopped, tick_progress)
}
VectorIndexEnum::HnswRam(index) => {
index.build_index_with_progress(permit, stopped, tick_progress)
}
VectorIndexEnum::HnswMmap(index) => {
index.build_index_with_progress(permit, stopped, tick_progress)
}
VectorIndexEnum::SparseRam(index) => {
index.build_index_with_progress(permit, stopped, tick_progress)
}
VectorIndexEnum::SparseMmap(index) => {
index.build_index_with_progress(permit, stopped, tick_progress)
}
}
}

17
lib/sparse/Cargo.toml
View File

@@ -8,6 +8,9 @@ authors = [
license = "Apache-2.0"
edition = "2021"
[features]
testing = []
[dependencies]
common = { path = "../common/common" }
io = { path = "../common/io" }
@@ -15,9 +18,23 @@ memory = { path = "../common/memory" }
memmap2 = "0.9.4"
schemars = { workspace = true }
serde = { workspace = true }
serde_json = { workspace = true }
tempfile = "3.10.1"
ordered-float = "4.2"
rand = "0.8.5"
validator = { workspace = true }
itertools = "0.12.1"
parking_lot = "0.12.2"
[dev-dependencies]
criterion = "0.5"
dataset = { path = "../common/dataset" }
indicatif = { workspace = true }
sparse = { path = ".", features = ["testing"] }
[target.'cfg(not(target_os = "windows"))'.dev-dependencies]
pprof = { workspace = true }
[[bench]]
name = "search"
harness = false

89
lib/sparse/benches/prof.rs Normal file
View File

@@ -0,0 +1,89 @@
use std::fs::File;
use std::io::Write;
use std::os::raw::c_int;
use std::path::Path;
use criterion::profiler::Profiler;
use pprof::flamegraph::TextTruncateDirection;
use pprof::protos::Message;
use pprof::ProfilerGuard;
/// 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 pprof_path = benchmark_dir.join("profile.pb");
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.hash = true;
options.image_width = Some(2500);
options.text_truncate_direction = TextTruncateDirection::Left;
options.font_size /= 3;
if let Some(profiler) = self.active_profiler.take() {
let report = profiler.report().build().unwrap();
let mut file = File::create(pprof_path).unwrap();
let profile = report.pprof().unwrap();
let mut content = Vec::new();
profile.encode(&mut content).unwrap();
file.write_all(&content).unwrap();
report
.flamegraph_with_options(flamegraph_file, &mut options)
.expect("Error writing flamegraph");
}
}
}

178
lib/sparse/benches/search.rs Normal file
View File

@@ -0,0 +1,178 @@
use std::io;
use std::path::Path;
use std::sync::atomic::AtomicBool;
use common::types::PointOffsetType;
use criterion::{criterion_group, criterion_main, Criterion};
use dataset::Dataset;
use indicatif::{ProgressBar, ProgressDrawTarget};
use itertools::Itertools;
use rand::rngs::StdRng;
use rand::SeedableRng as _;
use sparse::common::scores_memory_pool::ScoresMemoryPool;
use sparse::common::sparse_vector::SparseVector;
use sparse::common::sparse_vector_fixture::{random_positive_sparse_vector, random_sparse_vector};
use sparse::index::inverted_index::inverted_index_ram::InvertedIndexRam;
use sparse::index::inverted_index::inverted_index_ram_builder::InvertedIndexBuilder;
use sparse::index::loaders::{self, Csr};
use sparse::index::search_context::SearchContext;
mod prof;
const NUM_QUERIES: usize = 2048;
const MAX_SPARSE_DIM: usize = 30_000;
const TOP: usize = 10;
pub fn bench_search(c: &mut Criterion) {
bench_uniform_random(c, "random-50k", 50_000);
bench_uniform_random(c, "random-500k", 500_000);
let query_vectors =
loaders::load_csr_vecs(Dataset::NeurIps2023Queries.download().unwrap()).unwrap();
let index_1m = load_csr_index(Dataset::NeurIps2023_1M.download().unwrap(), 1.0).unwrap();
run_bench(c, "neurips2023-1M", index_1m, query_vectors.clone());
let index_full = load_csr_index(Dataset::NeurIps2023Full.download().unwrap(), 0.25).unwrap();
run_bench(c, "neurips2023-full-25pct", index_full, query_vectors);
bench_movies(c);
}
fn bench_uniform_random(c: &mut Criterion, name: &str, num_vectors: usize) {
let mut rnd = StdRng::seed_from_u64(42);
let index = InvertedIndexBuilder::build_from_iterator((0..num_vectors).map(|idx| {
(
idx as PointOffsetType,
random_sparse_vector(&mut rnd, MAX_SPARSE_DIM).into_remapped(),
)
}));
let query_vectors = (0..NUM_QUERIES)
.map(|_| random_positive_sparse_vector(&mut rnd, MAX_SPARSE_DIM))
.collect::<Vec<_>>();
run_bench(c, name, index, query_vectors);
}
pub fn bench_movies(c: &mut Criterion) {
let mut iter =
loaders::JsonReader::open(Dataset::SpladeWikiMovies.download().unwrap()).unwrap();
// Use the first NUM_QUERIES vectors as queries, and the rest as index.
let query_vectors = (0..NUM_QUERIES)
.map(|_| iter.next().unwrap().unwrap())
.collect_vec();
let index = InvertedIndexBuilder::build_from_iterator(
iter.enumerate()
.map(|(idx, vec)| (idx as PointOffsetType, vec.unwrap().into_remapped())),
);
run_bench(c, "movies", index, query_vectors);
}
pub fn run_bench(
c: &mut Criterion,
name: &str,
index: InvertedIndexRam,
mut query_vectors: Vec<SparseVector>,
) {
let pool = ScoresMemoryPool::new();
let stopped = AtomicBool::new(false);
let mut group = c.benchmark_group(format!("search/{}", name));
let mut it = query_vectors.iter().cycle();
group.bench_function("basic", |b| {
b.iter_batched(
|| it.next().unwrap().clone().into_remapped(),
|vec| SearchContext::new(vec, TOP, &index, pool.get(), &stopped).search(&|_| true),
criterion::BatchSize::SmallInput,
)
});
let hottest_id = index
.postings
.iter()
.enumerate()
.map(|(i, p)| (i, p.elements.len()))
.max_by_key(|(_, len)| *len)
.unwrap()
.0 as u32;
let average_elements = index
.postings
.iter()
.map(|p| p.elements.len())
.sum::<usize>() as f64
/ index.postings.len() as f64;
eprintln!(
"Hottest id: {hottest_id} (elements: {}), average elements: {average_elements}",
index.postings[hottest_id as usize].elements.len(),
);
for vec in &mut query_vectors {
vec.indices.truncate(4);
vec.values.truncate(4);
if let Err(idx) = vec.indices.binary_search(&hottest_id) {
if idx < vec.indices.len() {
vec.indices[idx] = hottest_id;
vec.values[idx] = 1.0;
} else {
vec.indices.push(hottest_id);
vec.values.push(1.0);
}
}
}
let mut it = query_vectors.iter().cycle();
group.bench_function("hottest", |b| {
b.iter(|| {
SearchContext::new(
it.next().unwrap().clone().into_remapped(),
TOP,
&index,
pool.get(),
&stopped,
)
.search(&|_| true)
})
});
}
fn load_csr_index(path: impl AsRef<Path>, ratio: f32) -> io::Result<InvertedIndexRam> {
let csr = Csr::open(path.as_ref())?;
let mut builder = InvertedIndexBuilder::new();
assert!(ratio > 0.0 && ratio <= 1.0);
let count = (csr.len() as f32 * ratio) as usize;
let bar =
ProgressBar::with_draw_target(Some(count as u64), ProgressDrawTarget::stderr_with_hz(12));
for (row, vec) in bar.wrap_iter(csr.iter().take(count).enumerate()) {
builder.add(
row as u32,
vec.map(|v| v.into_remapped())
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?,
);
}
bar.finish_and_clear();
Ok(builder.build())
}
#[cfg(not(target_os = "windows"))]
criterion_group! {
name = benches;
config = Criterion::default().with_profiler(prof::FlamegraphProfiler::new(100));
targets = bench_search,
}
#[cfg(target_os = "windows")]
criterion_group! {
name = benches;
config = Criterion::default();
targets = bench_search,
}
criterion_main!(benches);

9
lib/sparse/src/common/sparse_vector.rs
View File

@@ -207,6 +207,15 @@ impl SparseVector {
debug_assert!(result.validate().is_ok());
result
}
/// Create [RemappedSparseVector] from this vector in a naive way. Only suitable for testing.
#[cfg(feature = "testing")]
pub fn into_remapped(self) -> RemappedSparseVector {
RemappedSparseVector {
indices: self.indices,
values: self.values,
}
}
}
impl TryFrom<Vec<(u32, f32)>> for RemappedSparseVector {

11
lib/sparse/src/index/inverted_index/inverted_index_ram_builder.rs
View File

@@ -52,4 +52,15 @@ impl InvertedIndexBuilder {
vector_count,
}
}
/// Creates an [InvertedIndexRam] from an iterator of (id, vector) pairs.
pub fn build_from_iterator(
iter: impl Iterator<Item = (PointOffsetType, RemappedSparseVector)>,
) -> InvertedIndexRam {
let mut builder = InvertedIndexBuilder::new();
for (id, vector) in iter {
builder.add(id, vector);
}
builder.build()
}
}

157
lib/sparse/src/index/loaders.rs Normal file
View File

@@ -0,0 +1,157 @@
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, BufRead as _, BufReader, Lines};
use std::mem::size_of;
use std::path::Path;
use memmap2::Mmap;
use memory::mmap_ops::{open_read_mmap, transmute_from_u8, transmute_from_u8_to_slice};
use validator::ValidationErrors;
use crate::common::sparse_vector::SparseVector;
/// Compressed Sparse Row matrix, baked by memory-mapped file.
///
/// The layout of the memory-mapped file is as follows:
///
/// ======= =========== ========== ===================
/// name type size start
/// ======= =========== ========== ===================
/// nrow u64 8 0
/// ncol u64 8 8
/// nnz u64 8 16
/// indptr u64[nrow+1] 8*(nrow+1) 24
/// indices u32[nnz] 4*nnz 24+8*(nrow+1)
/// data u32[nnz] 4*nnz 24+8*(nrow+1)+4*nnz
pub struct Csr {
mmap: Mmap,
nrow: usize,
nnz: usize,
intptr: Vec<u64>,
}
const CSR_HEADER_SIZE: usize = size_of::<u64>() * 3;
impl Csr {
pub fn open(path: impl AsRef<Path>) -> io::Result<Self> {
Self::from_mmap(open_read_mmap(path.as_ref())?)
}
#[inline]
#[allow(clippy::len_without_is_empty)]
pub fn len(&self) -> usize {
self.nrow
}
pub fn iter(&self) -> CsrIter<'_> {
CsrIter { csr: self, row: 0 }
}
fn from_mmap(mmap: Mmap) -> io::Result<Self> {
let (nrow, ncol, nnz) =
transmute_from_u8::<(u64, u64, u64)>(&mmap.as_ref()[..CSR_HEADER_SIZE]);
let (nrow, _ncol, nnz) = (*nrow as usize, *ncol as usize, *nnz as usize);
let indptr = Vec::from(transmute_from_u8_to_slice::<u64>(
&mmap.as_ref()[CSR_HEADER_SIZE..CSR_HEADER_SIZE + size_of::<u64>() * (nrow + 1)],
));
if !indptr.windows(2).all(|w| w[0] <= w[1]) || indptr.last() != Some(&(nnz as u64)) {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"Invalid indptr array",
));
}
Ok(Self {
mmap,
nrow,
nnz,
intptr: indptr,
})
}
#[inline]
unsafe fn vec(&self, row: usize) -> Result<SparseVector, ValidationErrors> {
let start = *self.intptr.get_unchecked(row) as usize;
let end = *self.intptr.get_unchecked(row + 1) as usize;
let mut pos = CSR_HEADER_SIZE + size_of::<u64>() * (self.nrow + 1);
let indices = transmute_from_u8_to_slice::<u32>(
self.mmap
.as_ref()
.get_unchecked(pos + size_of::<u32>() * start..pos + size_of::<u32>() * end),
);
pos += size_of::<u32>() * self.nnz;
let data = transmute_from_u8_to_slice::<f32>(
self.mmap
.as_ref()
.get_unchecked(pos + size_of::<f32>() * start..pos + size_of::<f32>() * end),
);
SparseVector::new(indices.to_vec(), data.to_vec())
}
}
/// Iterator over the rows of a CSR matrix.
pub struct CsrIter<'a> {
csr: &'a Csr,
row: usize,
}
impl<'a> Iterator for CsrIter<'a> {
type Item = Result<SparseVector, ValidationErrors>;
fn next(&mut self) -> Option<Self::Item> {
(self.row < self.csr.nrow).then(|| {
let vec = unsafe { self.csr.vec(self.row) };
self.row += 1;
vec
})
}
}
impl<'a> ExactSizeIterator for CsrIter<'a> {
fn len(&self) -> usize {
self.csr.nrow - self.row
}
}
pub fn load_csr_vecs(path: impl AsRef<Path>) -> io::Result<Vec<SparseVector>> {
Csr::open(path)?
.iter()
.collect::<Result<Vec<_>, _>>()
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
}
/// Stream of sparse vectors in JSON format.
pub struct JsonReader(Lines<BufReader<File>>);
impl JsonReader {
pub fn open(path: impl AsRef<Path>) -> io::Result<Self> {
Ok(JsonReader(BufReader::new(File::open(path)?).lines()))
}
}
impl Iterator for JsonReader {
type Item = Result<SparseVector, io::Error>;
fn next(&mut self) -> Option<Self::Item> {
self.0.next().map(|line| {
line.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
.and_then(|line| {
let data: HashMap<String, f32> = serde_json::from_str(&line)?;
SparseVector::new(
data.keys()
.map(|k| k.parse())
.collect::<Result<Vec<_>, _>>()
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?,
data.values().copied().collect(),
)
.map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))
})
})
}
}

1
lib/sparse/src/index/mod.rs
View File

@@ -1,5 +1,6 @@
#![allow(dead_code)]
pub mod inverted_index;
pub mod loaders;
pub mod posting_list;
pub mod search_context;