use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use arrow_array::{Float64Array, StringArray, UInt32Array, UInt64Array};
use arrow_ipc::writer::StreamWriter;
use arrow_schema::Field;
use base64::Engine as _;
use base64::engine::general_purpose::STANDARD;
use futures::{StreamExt as _, stream};
use serde_json::json;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
use super::*;
use crate::scandal::SplitBatch;
use crate::scandal::api::{PlanScanRequest, ReadScanRequest, SearchVectorsRequest};
use crate::scandal::types::Scan as ScandalScan;
use crate::scandal::wire::{attach_progress, read_scan_schema};
use crate::vector_index::Binding;
use crate::{Client, RemoteScandal};
struct Fixture {
server: MockServer,
client: Client,
index: VectorIndex,
source: Table,
read: ReadVersionResponse,
}
impl Fixture {
async fn new() -> Self {
let server = MockServer::start().await;
let remote = RemoteScandal::new(server.uri().parse().unwrap()).unwrap();
let client = Client::builder()
.api_url(server.uri())
.bearer_token("caller")
.scandal(Arc::new(remote))
.build()
.unwrap();
let source = client.table_by_id(Uuid::new_v4(), Uuid::now_v7());
let index = VectorIndex {
project: client.project("demo").unwrap(),
binding: Binding::Commit,
resource_id: Uuid::new_v4(),
version_id: Uuid::now_v7(),
};
let read = serde_json::from_value(json!({
"definition": {
"sourceResourceId": source.identity().resource_id,
"sourceTableId": source.identity().resource_id,
"sourceFieldId": Uuid::new_v4(),
"dimensions": 2, "metric": "cosine", "engineId": "vortex.ivf-rabitq", "engineOptions": {},
},
"rootUri": "spfsu://index/",
"version": {
"vectorIndexId": index.resource_id, "versionId": index.version_id,
"createdAtMicros": "0", "sealedAtMicros": "0",
"sourceVersionId": source.identity().version_id,
"physical": {
"codec": "vortex.ivf-rabitq", "codecVersion": 3, "indexedRows": "3",
"manifestPath": "manifest.vortex", "partitions": 2,
"metadata": STANDARD.encode(format!("32:{}", "a".repeat(64))),
},
},
})).unwrap();
Self {
server,
client,
index,
source,
read,
}
}
async fn mount_read(&self) {
Mock::given(method("GET"))
.and(path(format!(
"/api/vector-indexes/{}/versions/{}",
self.index.resource_id, self.index.version_id
)))
.respond_with(ResponseTemplate::new(200).set_body_json(&self.read))
.expect(1)
.mount(&self.server)
.await;
}
async fn mount_bindings(&self) {
let expiry = (SystemTime::now().duration_since(UNIX_EPOCH).unwrap()
+ Duration::from_mins(15))
.as_secs();
for (resource, root, token) in [
(
format!("vector-indexes/{}", self.index.resource_id),
"spfsu://index-version/",
"index-grant",
),
(
format!("tables/{}", self.source.identity().resource_id),
"spfsu://source/",
"source-grant",
),
] {
Mock::given(method("POST"))
.and(path(format!("/api/{resource}/storage-binding")))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({
"rootUri": root, "mode": "read_only", "expiresAt": expiry.to_string(),
"access": { "type": "signed", "signed": { "endpoint": "", "token": token } },
})))
.expect(1)
.mount(&self.server)
.await;
}
Mock::given(method("GET"))
.and(path(format!(
"/api/tables/{}/versions/{}",
self.source.identity().resource_id,
self.source.identity().version_id
)))
.respond_with(
ResponseTemplate::new(200)
.set_body_json(json!({ "rootUri": "spfsu://source/", "seqNo": "1" })),
)
.expect(1)
.mount(&self.server)
.await;
}
fn artifact(&self) -> VectorIndexArtifact {
self.index
.artifact(self.read.clone(), &self.source)
.unwrap()
}
}
fn neighbors(artifact: &VectorIndexArtifact, ids: Vec<u64>) -> RecordBatch {
let info = VectorSearchInfo {
origin: artifact.origin.clone(),
metric: VectorMetric::Cosine,
exhaustive: true,
indexed_rows: 3.into(),
limit: 3,
probes: 2,
rerank: 3,
};
let schema = Arc::new(Schema::new_with_metadata(
vec![
Field::new(ROW_ID_COLUMN, DataType::UInt64, false),
Field::new(DISTANCE_COLUMN, DataType::Float64, false),
Field::new(RANK_COLUMN, DataType::UInt32, false),
],
[(INFO_KEY.into(), serde_json::to_string(&info).unwrap())].into(),
));
let count = ids.len();
RecordBatch::try_new(
schema,
vec![
Arc::new(UInt64Array::from(ids)),
Arc::new(Float64Array::from(
(0..count)
.map(|i| (i + 1) as f64 / 10.0)
.collect::<Vec<_>>(),
)),
Arc::new(UInt32Array::from((1..=count as u32).collect::<Vec<_>>())),
],
)
.unwrap()
}
fn ipc(batch: &RecordBatch) -> Vec<u8> {
let mut bytes = Vec::new();
{
let mut writer = StreamWriter::try_new(&mut bytes, &batch.schema()).unwrap();
writer.write(batch).unwrap();
writer.finish().unwrap();
}
bytes
}
fn source_rows(scan: &ScandalScan) -> RecordBatch {
assert_eq!(scan.columns.as_deref().unwrap(), ["label"]);
assert_eq!(scan.row_id_column.as_deref(), Some(ROW_ID_COLUMN));
let ids = [10u64, 20, 30]
.into_iter()
.filter(|id| {
scan.row_ids
.as_ref()
.unwrap()
.iter()
.any(|range| (range.start.0..range.end.0).contains(id))
})
.collect::<Vec<_>>();
RecordBatch::try_new(
Arc::new(Schema::new(vec![
Field::new("label", DataType::Utf8, false),
Field::new(ROW_ID_COLUMN, DataType::UInt64, false),
])),
vec![
Arc::new(StringArray::from(
ids.iter().map(|id| format!("row-{id}")).collect::<Vec<_>>(),
)),
Arc::new(UInt64Array::from(ids)),
],
)
.unwrap()
}
#[tokio::test]
async fn pages_preserve_rank_without_repeating_the_search() {
let fixture = Fixture::new().await;
fixture.mount_read().await;
fixture.mount_bindings().await;
let artifact = fixture.artifact();
Mock::given(method("POST"))
.and(path("/vector-indexes/search"))
.respond_with(move |request: &wiremock::Request| {
let request: SearchVectorsRequest = serde_json::from_slice(&request.body).unwrap();
let value = json!(&request);
assert!(value.get("source").is_none());
assert!(value["index"]["origin"].get("sourceTableId").is_none());
assert_eq!(request.index, artifact);
assert_eq!(request.source_storage.credential(), Some("source-grant"));
assert_eq!(request.index_storage.credential(), Some("index-grant"));
assert!(request.columns.is_none());
assert_eq!(request.exclude_row_ids.unwrap(), [999u64.into()]);
ResponseTemplate::new(200)
.insert_header("content-type", "application/vnd.apache.arrow.stream")
.set_body_bytes(ipc(&neighbors(&artifact, vec![30, 10, 20])))
})
.expect(1)
.mount(&fixture.server)
.await;
Mock::given(method("POST")).and(path("/scans/plan"))
.respond_with(|request: &wiremock::Request| {
let request: PlanScanRequest = serde_json::from_slice(&request.body).unwrap();
let rows = source_rows(&request.scan);
ResponseTemplate::new(200).set_body_json(json!({
"info": { "schema": rows.schema(), "filterClauses": [] },
"splits": [{ "split": { "fragment": "0", "rows": { "start": "0", "end": rows.num_rows().to_string() } },
"estimatedRows": rows.num_rows().to_string() }],
}))
}).expect(2).mount(&fixture.server).await;
Mock::given(method("POST"))
.and(path("/scans/read"))
.respond_with(|request: &wiremock::Request| {
let request: ReadScanRequest = serde_json::from_slice(&request.body).unwrap();
let batch = source_rows(&request.scan);
let schema = read_scan_schema(&batch.schema());
let batch = attach_progress(
SplitBatch {
batch,
next_row: request.split.rows.end.0,
metrics: None,
},
&schema,
)
.unwrap();
ResponseTemplate::new(200)
.insert_header("content-type", "application/vnd.apache.arrow.stream")
.set_body_bytes(ipc(&batch))
})
.expect(2)
.mount(&fixture.server)
.await;
let results = fixture
.index
.search(
&fixture.source,
vec![1.0, 0.0],
VectorSearchOptions {
limit: 3,
probes: Some(2),
rerank: Some(3),
exclude_row_ids: vec![999],
..Default::default()
},
)
.await
.unwrap();
assert_eq!(results.query(), [1.0, 0.0]);
assert!(results.info().exhaustive);
let first = results.page(0..2, vec!["label".into()]).await.unwrap();
let last = results.page(2..3, vec!["label".into()]).await.unwrap();
let labels = |batch: &RecordBatch| {
batch
.column_by_name("label")
.unwrap()
.as_any()
.downcast_ref::<StringArray>()
.unwrap()
.iter()
.map(|value| value.unwrap().to_owned())
.collect::<Vec<_>>()
};
assert_eq!(labels(&first), ["row-30", "row-10"]);
assert_eq!(labels(&last), ["row-20"]);
assert_eq!(
results.page(0..1, Vec::new()).await.unwrap().num_columns(),
3
);
assert!(results.page(0..4, Vec::new()).await.is_err());
fixture.server.verify().await;
}
#[tokio::test]
async fn a_mismatched_source_fails_before_storage_access() {
for wrong_identity in [true, false] {
let fixture = Fixture::new().await;
fixture.mount_read().await;
let exact = fixture.source.identity();
let (resource_id, version_id) = if wrong_identity {
(Uuid::new_v4(), exact.version_id)
} else {
(exact.resource_id, Uuid::now_v7())
};
let source = fixture.client.table_by_id(resource_id, version_id);
let error = fixture
.index
.search(&source, vec![1.0, 0.0], VectorSearchOptions::default())
.await
.err()
.unwrap();
assert!(
error.to_string().contains("source does not match"),
"{error}"
);
assert_eq!(fixture.server.received_requests().await.unwrap().len(), 1);
}
}
#[tokio::test]
async fn result_collection_enforces_the_fixed_limit() {
let fixture = Fixture::new().await;
let artifact = fixture.artifact();
let batch = neighbors(&artifact, vec![10, 20, 30, 40]);
let schema = batch.schema();
let evaluation = Evaluation {
schema,
batches: stream::iter([Ok(batch)]).boxed(),
};
let error = VectorSearch::read(
fixture.source,
vec![1.0, 0.0],
VectorSearchOptions {
limit: 3,
..Default::default()
},
&artifact,
evaluation,
)
.await
.err()
.unwrap();
assert!(error.to_string().contains("row limit"), "{error}");
}