feat(rust): add with_row_id to rust SDK (#1683)

This commit is contained in:
Lei Xu
2024-09-21 21:26:19 -07:00
committed by GitHub
parent 11072b9edc
commit 7a8d2f37c4
2 changed files with 38 additions and 0 deletions

View File

@@ -402,6 +402,9 @@ pub trait QueryBase {
///
/// By default, it is false.
fn fast_search(self) -> Self;
/// Return the `_rowid` meta column from the Table.
fn with_row_id(self) -> Self;
}
pub trait HasQuery {
@@ -438,6 +441,11 @@ impl<T: HasQuery> QueryBase for T {
self.mut_query().fast_search = true;
self
}
fn with_row_id(mut self) -> Self {
self.mut_query().with_row_id = true;
self
}
}
/// Options for controlling the execution of a query
@@ -548,6 +556,11 @@ pub struct Query {
///
/// By default, this is false.
pub(crate) fast_search: bool,
/// If set to true, the query will return the `_rowid` meta column.
///
/// By default, this is false.
pub(crate) with_row_id: bool,
}
impl Query {
@@ -560,6 +573,7 @@ impl Query {
full_text_search: None,
select: Select::All,
fast_search: false,
with_row_id: false,
}
}
@@ -1160,4 +1174,24 @@ mod tests {
.unwrap();
assert!(!plan.contains("Take"));
}
#[tokio::test]
async fn test_with_row_id() {
let tmp_dir = tempdir().unwrap();
let table = make_test_table(&tmp_dir).await;
let results = table
.vector_search(&[0.1, 0.2, 0.3, 0.4])
.unwrap()
.with_row_id()
.limit(10)
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
for batch in results {
assert!(batch.column_by_name("_rowid").is_some());
}
}
}

View File

@@ -1938,6 +1938,10 @@ impl TableInternal for NativeTable {
Select::All => {}
}
if query.base.with_row_id {
scanner.with_row_id();
}
if let Some(opts) = options {
scanner.batch_size(opts.max_batch_length as usize);
}