mirror of
https://github.com/hanzoai/vector.git
synced 2026-10-30 18:18:56 +00:00
edd6be8a90
* feat: add shard-ops timeout * feat: add shard-ops timeout * feat: add missing timeout * fix: format * fix: type serialization * fix: bad wal flush logic * feat: add seperate wait_time * fix: linter issues * fix: linter * chore: make select more readable * fix: linter issues * chore: remove comment * chore: adjust comment * chore: update definition * chore: remove extra wait_timeout -> make timeout global * chore: remove wait_timeout from point_ops * chore: remove wait_timeout from grpc * refactor: make shard_ops more readable * chore: remove timeout from update_all * feat: propagate wait_timeout error * fix: clippy recommendation * chore: update wait timeout * refactor: update merge_successful_update_results * fix: tests inconsistency * fix: revert global ord, partialord for UpdateStatus * fix: integration test * fix: res ingoring * fix: review requests * fix: clippy * chore: revert back futures unordered * fix: resolve response status * force timeout for remote operations * fix: pr reviews * fix: switch to duration * Bump OpenAPI spec * feat: add serde_as * fix: timeout statu condition * Minor tweaks * Use tokio timeout directly rather than deadline * Revert schema generator permissions * Remove obsolete semi-colon --------- Co-authored-by: generall <[email protected]> Co-authored-by: timvisee <[email protected]>
307 lines
9.3 KiB
Rust
307 lines
9.3 KiB
Rust
use collection::operations::CollectionUpdateOperations;
|
|
use collection::operations::point_ops::{
|
|
BatchPersisted, BatchVectorStructPersisted, PointInsertOperationsInternal, PointOperations,
|
|
WriteOrdering,
|
|
};
|
|
use collection::operations::shard_selector_internal::ShardSelectorInternal;
|
|
use collection::operations::types::ScrollRequestInternal;
|
|
use common::counter::hardware_accumulator::HwMeasurementAcc;
|
|
use itertools::Itertools;
|
|
use segment::json_path::JsonPath;
|
|
use segment::types::{PayloadContainer, PayloadSelectorExclude, WithPayloadInterface};
|
|
use serde_json::Value;
|
|
use tempfile::Builder;
|
|
|
|
use crate::common::{N_SHARDS, load_local_collection, simple_collection_fixture};
|
|
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn test_collection_reloading() {
|
|
test_collection_reloading_with_shards(1).await;
|
|
test_collection_reloading_with_shards(N_SHARDS).await;
|
|
}
|
|
|
|
async fn test_collection_reloading_with_shards(shard_number: u32) {
|
|
let collection_dir = Builder::new().prefix("collection").tempdir().unwrap();
|
|
|
|
let collection = simple_collection_fixture(collection_dir.path(), shard_number).await;
|
|
|
|
collection.stop_gracefully().await;
|
|
|
|
for _i in 0..5 {
|
|
let collection_path = collection_dir.path();
|
|
let collection = load_local_collection(
|
|
"test".to_string(),
|
|
collection_path,
|
|
&collection_path.join("snapshots"),
|
|
)
|
|
.await;
|
|
let insert_points =
|
|
CollectionUpdateOperations::PointOperation(PointOperations::UpsertPoints(
|
|
PointInsertOperationsInternal::PointsBatch(BatchPersisted {
|
|
ids: vec![0, 1].into_iter().map(|x| x.into()).collect_vec(),
|
|
vectors: BatchVectorStructPersisted::Single(vec![
|
|
vec![1.0, 0.0, 1.0, 1.0],
|
|
vec![1.0, 0.0, 1.0, 0.0],
|
|
]),
|
|
payloads: None,
|
|
}),
|
|
));
|
|
let hw_counter = HwMeasurementAcc::new();
|
|
collection
|
|
.update_from_client_simple(
|
|
insert_points,
|
|
true,
|
|
None,
|
|
WriteOrdering::default(),
|
|
hw_counter,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
collection.stop_gracefully().await;
|
|
}
|
|
|
|
let collection_path = collection_dir.path();
|
|
let collection = load_local_collection(
|
|
"test".to_string(),
|
|
collection_path,
|
|
&collection_path.join("snapshots"),
|
|
)
|
|
.await;
|
|
assert_eq!(
|
|
collection
|
|
.info(&ShardSelectorInternal::All)
|
|
.await
|
|
.unwrap()
|
|
.points_count,
|
|
Some(2),
|
|
);
|
|
|
|
collection.stop_gracefully().await;
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn test_collection_payload_reloading() {
|
|
test_collection_payload_reloading_with_shards(1).await;
|
|
test_collection_payload_reloading_with_shards(N_SHARDS).await;
|
|
}
|
|
|
|
async fn test_collection_payload_reloading_with_shards(shard_number: u32) {
|
|
let collection_dir = Builder::new().prefix("collection").tempdir().unwrap();
|
|
{
|
|
let collection = simple_collection_fixture(collection_dir.path(), shard_number).await;
|
|
let insert_points =
|
|
CollectionUpdateOperations::PointOperation(PointOperations::UpsertPoints(
|
|
PointInsertOperationsInternal::PointsBatch(BatchPersisted {
|
|
ids: vec![0, 1].into_iter().map(|x| x.into()).collect_vec(),
|
|
vectors: BatchVectorStructPersisted::Single(vec![
|
|
vec![1.0, 0.0, 1.0, 1.0],
|
|
vec![1.0, 0.0, 1.0, 0.0],
|
|
]),
|
|
payloads: serde_json::from_str(r#"[{ "k": "v1" } , { "k": "v2"}]"#).unwrap(),
|
|
}),
|
|
));
|
|
let hw_counter = HwMeasurementAcc::new();
|
|
collection
|
|
.update_from_client_simple(
|
|
insert_points,
|
|
true,
|
|
None,
|
|
WriteOrdering::default(),
|
|
hw_counter,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
collection.stop_gracefully().await;
|
|
}
|
|
let collection_path = collection_dir.path();
|
|
let collection = load_local_collection(
|
|
"test".to_string(),
|
|
collection_path,
|
|
&collection_path.join("snapshots"),
|
|
)
|
|
.await;
|
|
|
|
let res = collection
|
|
.scroll_by(
|
|
ScrollRequestInternal {
|
|
offset: None,
|
|
limit: Some(10),
|
|
filter: None,
|
|
with_payload: Some(WithPayloadInterface::Bool(true)),
|
|
with_vector: true.into(),
|
|
order_by: None,
|
|
},
|
|
None,
|
|
&ShardSelectorInternal::All,
|
|
None,
|
|
HwMeasurementAcc::new(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(res.points.len(), 2);
|
|
|
|
match res.points[0]
|
|
.payload
|
|
.as_ref()
|
|
.expect("has payload")
|
|
.get_value(&JsonPath::new("k"))
|
|
.into_iter()
|
|
.next()
|
|
.expect("has value")
|
|
{
|
|
Value::String(value) => assert_eq!("v1", value),
|
|
_ => panic!("unexpected type"),
|
|
}
|
|
|
|
eprintln!(
|
|
"res = {:#?}",
|
|
res.points[0]
|
|
.payload
|
|
.as_ref()
|
|
.unwrap()
|
|
.get_value(&JsonPath::new("k"))
|
|
);
|
|
|
|
collection.stop_gracefully().await;
|
|
}
|
|
|
|
#[tokio::test(flavor = "multi_thread")]
|
|
async fn test_collection_payload_custom_payload() {
|
|
test_collection_payload_custom_payload_with_shards(1).await;
|
|
test_collection_payload_custom_payload_with_shards(N_SHARDS).await;
|
|
}
|
|
|
|
async fn test_collection_payload_custom_payload_with_shards(shard_number: u32) {
|
|
let collection_dir = Builder::new().prefix("collection").tempdir().unwrap();
|
|
{
|
|
let collection = simple_collection_fixture(collection_dir.path(), shard_number).await;
|
|
let insert_points =
|
|
CollectionUpdateOperations::PointOperation(PointOperations::UpsertPoints(
|
|
PointInsertOperationsInternal::PointsBatch(BatchPersisted {
|
|
ids: vec![0.into(), 1.into()],
|
|
vectors: BatchVectorStructPersisted::Single(vec![
|
|
vec![1.0, 0.0, 1.0, 1.0],
|
|
vec![1.0, 0.0, 1.0, 0.0],
|
|
]),
|
|
payloads: serde_json::from_str(
|
|
r#"[{ "k1": "v1" }, { "k1": "v2" , "k2": "v3", "k3": "v4"}]"#,
|
|
)
|
|
.unwrap(),
|
|
}),
|
|
));
|
|
let hw_counter = HwMeasurementAcc::new();
|
|
collection
|
|
.update_from_client_simple(
|
|
insert_points,
|
|
true,
|
|
None,
|
|
WriteOrdering::default(),
|
|
hw_counter,
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
collection.stop_gracefully().await;
|
|
}
|
|
|
|
let collection_path = collection_dir.path();
|
|
let collection = load_local_collection(
|
|
"test".to_string(),
|
|
collection_path,
|
|
&collection_path.join("snapshots"),
|
|
)
|
|
.await;
|
|
|
|
// Test res with filter payload
|
|
let res_with_custom_payload = collection
|
|
.scroll_by(
|
|
ScrollRequestInternal {
|
|
offset: None,
|
|
limit: Some(10),
|
|
filter: None,
|
|
with_payload: Some(WithPayloadInterface::Fields(vec![JsonPath::new("k2")])),
|
|
with_vector: true.into(),
|
|
order_by: None,
|
|
},
|
|
None,
|
|
&ShardSelectorInternal::All,
|
|
None,
|
|
HwMeasurementAcc::new(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert!(
|
|
res_with_custom_payload.points[0]
|
|
.payload
|
|
.as_ref()
|
|
.expect("has payload")
|
|
.is_empty(),
|
|
);
|
|
|
|
match res_with_custom_payload.points[1]
|
|
.payload
|
|
.as_ref()
|
|
.expect("has payload")
|
|
.get_value(&JsonPath::new("k2"))
|
|
.into_iter()
|
|
.next()
|
|
.expect("has value")
|
|
{
|
|
Value::String(value) => assert_eq!("v3", value),
|
|
_ => panic!("unexpected type"),
|
|
}
|
|
|
|
// Test res with filter payload dict
|
|
let res_with_custom_payload = collection
|
|
.scroll_by(
|
|
ScrollRequestInternal {
|
|
offset: None,
|
|
limit: Some(10),
|
|
filter: None,
|
|
with_payload: Some(PayloadSelectorExclude::new(vec![JsonPath::new("k1")]).into()),
|
|
with_vector: false.into(),
|
|
order_by: None,
|
|
},
|
|
None,
|
|
&ShardSelectorInternal::All,
|
|
None,
|
|
HwMeasurementAcc::new(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
assert!(
|
|
res_with_custom_payload.points[0]
|
|
.payload
|
|
.as_ref()
|
|
.expect("has payload")
|
|
.is_empty(),
|
|
);
|
|
|
|
assert_eq!(
|
|
res_with_custom_payload.points[1]
|
|
.payload
|
|
.as_ref()
|
|
.expect("has payload")
|
|
.len(),
|
|
2
|
|
);
|
|
|
|
match res_with_custom_payload.points[1]
|
|
.payload
|
|
.as_ref()
|
|
.expect("has payload")
|
|
.get_value(&JsonPath::new("k3"))
|
|
.into_iter()
|
|
.next()
|
|
.expect("has value")
|
|
{
|
|
Value::String(value) => assert_eq!("v4", value),
|
|
_ => panic!("unexpected type"),
|
|
};
|
|
|
|
collection.stop_gracefully().await;
|
|
}
|