feat: LocalTable for vectordb now supports filters without vector search (#693)

Note this currently the filter/where is only implemented for LocalTable
so that it requires an explicit cast to "enable" (see new unit test).
The alternative is to add it to the Table interface, but since it's not
available on RemoteTable this may cause some user experience issues.
This commit is contained in:
Chang She
2023-12-13 22:59:01 -08:00
committed by GitHub
parent 63ee8fa6a1
commit 098e397cf0
7 changed files with 104 additions and 37 deletions

View File

@@ -23,8 +23,14 @@ impl JsQuery {
let query_obj = cx.argument::<JsObject>(0)?;
let limit = query_obj
.get::<JsNumber, _, _>(&mut cx, "_limit")?
.value(&mut cx);
.get_opt::<JsNumber, _, _>(&mut cx, "_limit")?
.map(|value| {
let limit = value.value(&mut cx) as u64;
if limit <= 0 {
panic!("Limit must be a positive integer");
}
limit
});
let select = query_obj
.get_opt::<JsArray, _, _>(&mut cx, "_select")?
.map(|arr| {
@@ -61,20 +67,23 @@ impl JsQuery {
let (deferred, promise) = cx.promise();
let channel = cx.channel();
let query_vector = query_obj.get::<JsArray, _, _>(&mut cx, "_queryVector")?;
let query = convert::js_array_to_vec(query_vector.deref(), &mut cx);
let query_vector = query_obj.get_opt::<JsArray, _, _>(&mut cx, "_queryVector")?;
let table = js_table.table.clone();
let query = query_vector.map(|q| convert::js_array_to_vec(q.deref(), &mut cx));
rt.spawn(async move {
let builder = table
.search(Float32Array::from(query))
.limit(limit as usize)
let mut builder = table
.search(query.map(|q| Float32Array::from(q)))
.refine_factor(refine_factor)
.nprobes(nprobes)
.filter(filter)
.metric_type(metric_type)
.select(select)
.prefilter(prefilter);
if let Some(limit) = limit {
builder = builder.limit(limit as usize);
};
let record_batch_stream = builder.execute();
let results = record_batch_stream
.and_then(|stream| {