rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-10-26 13:36:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat!: return lazy iterators for table and namespace listings (#4599)

Browse Source 
Make table and namespace listings consistent with other resource
listings across Rust, Python, and TypeScript. Callers can iterate over
individual names across all pages, with configurable page sizes and
continuation tokens.

In order to avoid breaking existing Python users, this change adds some
compatibility methods to the returned Python iterators. Code such as:
```
for t in db.list_tables().tables:
    print(t)
```
will continue to work as expected because the iterator has a `tables`
field and a `page_token` field, which function as expected to reveal the
first page of results. In other words, it is never necessary to use the
iterator __as an iterator__. It may be used as if it were an old-style
response object. (As the iterator progresses, the `tables` list may
switch to the second, third, etc. page of results.) Similar
compatibility code allows the async python API to keep working as well.
This commit is contained in:
Colin Patrick McCabe
2026-10-22 15:41:30 -07:00
committed by GitHub
parent 88203d2515
commit 631ea41a63
27 changed files with 814 additions and 713 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/lancedb/src/connection.rs
+90 -71
View File
@@ -13,7 +13,7 @@ use lance::dataset::refs::MAIN_BRANCH;
use lance_namespace::models::{
CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest,
DescribeNamespaceResponse, DropNamespaceRequest, DropNamespaceResponse, ListNamespacesRequest,
ListNamespacesResponse, ListTablesRequest, ListTablesResponse,
ListTablesRequest,
};
#[cfg(feature = "aws")]
use object_store::aws::AwsCredential;
@@ -28,7 +28,7 @@ use crate::database::{
};
use crate::embeddings::{EmbeddingRegistry, MemoryRegistry};
use crate::error::{Error, Result};
use crate::listing::{Listing, ListingOptions};
use crate::listing::{Listing, ListingOptions, ListingPage};
#[cfg(feature = "remote")]
use crate::remote::{
client::ClientConfig,
@@ -1126,12 +1126,43 @@ impl Connection {
self.internal.drop_all_tables(namespace_path).await
}
/// List immediate child namespace names in the given namespace
pub async fn list_namespaces(
/// Lazily list immediate child namespace names. Each request asks for at
/// most 100 names unless `options.page_limit` is set.
///
/// ```
/// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> {
/// use futures::TryStreamExt;
/// let names: Vec<String> = connection.list_namespaces(&[], Default::default())
/// .try_collect().await?;
/// # Ok(())
/// # }
/// ```
pub fn list_namespaces(
&self,
request: ListNamespacesRequest,
) -> Result<ListNamespacesResponse> {
self.internal.list_namespaces(request).await
namespace_path: &[String],
options: ListingOptions,
) -> Listing<String> {
let database = self.database().clone();
let namespace_path = namespace_path.to_vec();
Listing::new(options, move |options| {
let database = database.clone();
let namespace_path = namespace_path.clone();
async move {
validate_namespace(&namespace_path)?;
let response = database
.list_namespaces(ListNamespacesRequest {
id: Some(namespace_path),
page_token: options.page_token,
limit: Some(options.page_limit.unwrap_or(100) as i32),
..Default::default()
})
.await?;
Ok(ListingPage {
items: response.namespaces,
page_token: response.page_token,
})
}
})
}
/// Create a new namespace
@@ -1182,23 +1213,43 @@ impl Connection {
self.internal.namespace_client_config().await
}
/// List tables with pagination support
/// Lazily list tables. Each request asks for at most 100 names unless
/// `options.page_limit` is set.
///
/// The default limit is 100 tables per page for both local and remote connections,
/// including namespaces. Zero returns an empty page without a continuation token.
/// Follow the response's opaque `page_token` until it is absent to retrieve every table;
/// a page can contain fewer than the limit even when more tables remain.
pub async fn list_tables(&self, mut request: ListTablesRequest) -> Result<ListTablesResponse> {
// Apply the SDK's page size here so every backend receives the same request.
let limit = request.limit.get_or_insert(100);
if *limit <= 0 {
return Ok(ListTablesResponse {
context: None,
tables: Vec::new(),
page_token: None,
});
}
self.internal.list_tables(request).await
/// ```
/// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> {
/// use futures::TryStreamExt;
/// let names: Vec<String> = connection.list_tables(&[], Default::default())
/// .try_collect().await?;
/// # Ok(())
/// # }
/// ```
pub fn list_tables(
&self,
namespace_path: &[String],
options: ListingOptions,
) -> Listing<String> {
let database = self.database().clone();
let namespace_path = namespace_path.to_vec();
Listing::new(options, move |options| {
let database = database.clone();
let namespace_path = namespace_path.clone();
async move {
validate_namespace(&namespace_path)?;
let response = database
.list_tables(ListTablesRequest {
id: Some(namespace_path),
page_token: options.page_token,
limit: Some(options.page_limit.unwrap_or(100) as i32),
..Default::default()
})
.await?;
Ok(ListingPage {
items: response.tables,
page_token: response.page_token,
})
}
})
}
/// Get the in-memory embedding registry.
@@ -2260,26 +2311,11 @@ mod tests {
}
names.sort();
// Walking in pages has to reach every table exactly once, with nothing lost at a
// page boundary.
let mut seen = Vec::with_capacity(names.len());
let mut page_token = None;
loop {
let page = db
.list_tables(ListTablesRequest {
id: Some(Vec::new()),
limit: Some(2),
page_token,
..Default::default()
})
.await
.unwrap();
seen.extend(page.tables);
page_token = page.page_token.filter(|token| !token.is_empty());
if page_token.is_none() {
break;
}
}
let seen = futures::TryStreamExt::try_collect::<Vec<_>>(
db.list_tables(&[], ListingOptions::default().page_limit(2)),
)
.await
.unwrap();
assert_eq!(seen, names);
}
@@ -2305,37 +2341,20 @@ mod tests {
.is_empty()
);
let first = db.list_tables(ListTablesRequest::default()).await.unwrap();
assert_eq!(first.tables, names[..100]);
let mut listing = db.list_tables(&[], ListingOptions::default());
assert_eq!(listing.num_page_results(), 0);
let first = listing.next_page().await.unwrap().unwrap();
assert_eq!(first.items, names[..100]);
assert!(first.page_token.is_some());
let second = db
.list_tables(ListTablesRequest {
page_token: first.page_token,
..Default::default()
})
.await
.unwrap();
assert_eq!(second.tables, names[100..]);
let second = listing.next_page().await.unwrap().unwrap();
assert_eq!(second.items, names[100..]);
assert!(second.page_token.is_none());
assert!(listing.next_page().await.unwrap().is_none());
let all = db
.list_tables(ListTablesRequest {
limit: Some(200),
..Default::default()
})
.await
.unwrap();
assert_eq!(all.tables, names);
assert!(all.page_token.is_none());
let empty = db
.list_tables(ListTablesRequest {
limit: Some(0),
..Default::default()
})
.await
.unwrap();
assert!(empty.tables.is_empty());
assert!(empty.page_token.is_none());
let mut all = db.list_tables(&[], ListingOptions::default().page_limit(200));
assert_eq!(all.next_page().await.unwrap().unwrap().items, names);
let mut invalid = db.list_tables(&[], ListingOptions::default().page_limit(0));
assert!(invalid.next_page().await.is_err());
}
#[tokio::test]
rust/lancedb/src/listing.rs
+93 -37
View File
@@ -51,8 +51,9 @@ type FetchPage<T> =
/// A lazy stream of resources with inspectable pagination state.
///
/// No requests are made until the stream is polled. Empty pages with continuation
/// tokens are skipped. Errors terminate the stream and retain the failed request's
/// No requests are made until the stream is polled or [`Self::next_page`] is called.
/// Item iteration skips empty pages with continuation tokens; `next_page` returns
/// them. Errors terminate the listing and retain the failed request's
/// token, allowing a new stream to retry it. Repeated tokens are rejected.
///
/// A token advances when its page is fetched. Drain [`Self::num_page_results`]
@@ -100,6 +101,64 @@ impl<T> Listing<T> {
pub fn page_token(&self) -> Option<&str> {
self.options.page_token.as_deref()
}
/// Return exactly one page without consuming a following page.
///
/// If item iteration has started, this returns the remaining items in its
/// cached page. An empty page with a continuation token is returned as is.
pub async fn next_page(&mut self) -> Result<Option<ListingPage<T>>> {
if self.items.len() == 0 {
if self.done {
return Ok(None);
}
futures::future::poll_fn(|cx| self.poll_fetch_page(cx)).await?;
}
let items = std::mem::replace(&mut self.items, Vec::new().into_iter()).collect();
Ok(Some(ListingPage {
items,
page_token: self.options.page_token.clone(),
}))
}
fn poll_fetch_page(&mut self, cx: &mut Context<'_>) -> Poll<Result<()>> {
if self
.options
.page_limit
.is_some_and(|limit| limit == 0 || limit > i32::MAX as u32)
{
self.done = true;
return Poll::Ready(Err(Error::InvalidInput {
message: "Listing page limit must be between 1 and 2147483647".into(),
}));
}
let pending = self
.pending
.get_or_insert_with(|| (self.fetch)(self.options.clone()));
let response = futures::ready!(pending.as_mut().poll(cx));
self.pending = None;
match response {
Ok(response) => {
let token = response.page_token.filter(|token| !token.is_empty());
if token
.as_ref()
.is_some_and(|token| !self.seen_tokens.insert(token.clone()))
{
self.done = true;
return Poll::Ready(Err(Error::Runtime {
message: "Listing response repeated a page_token".into(),
}));
}
self.done = token.is_none();
self.options.page_token = token;
self.items = response.items.into_iter();
Poll::Ready(Ok(()))
}
Err(error) => {
self.done = true;
Poll::Ready(Err(error))
}
}
}
}
impl<T: Unpin> Stream for Listing<T> {
@@ -114,41 +173,8 @@ impl<T: Unpin> Stream for Listing<T> {
if this.done {
return Poll::Ready(None);
}
if this
.options
.page_limit
.is_some_and(|limit| limit == 0 || limit > i32::MAX as u32)
{
this.done = true;
return Poll::Ready(Some(Err(Error::InvalidInput {
message: "Listing page limit must be between 1 and 2147483647".into(),
})));
}
let pending = this
.pending
.get_or_insert_with(|| (this.fetch)(this.options.clone()));
let response = futures::ready!(pending.as_mut().poll(cx));
this.pending = None;
match response {
Ok(response) => {
let token = response.page_token.filter(|token| !token.is_empty());
if token
.as_ref()
.is_some_and(|token| !this.seen_tokens.insert(token.clone()))
{
this.done = true;
return Poll::Ready(Some(Err(Error::Runtime {
message: "Listing response repeated a page_token".into(),
})));
}
this.done = token.is_none();
this.options.page_token = token;
this.items = response.items.into_iter();
}
Err(error) => {
this.done = true;
return Poll::Ready(Some(Err(error)));
}
if let Err(error) = futures::ready!(this.poll_fetch_page(cx)) {
return Poll::Ready(Some(Err(error)));
}
}
}
@@ -224,6 +250,36 @@ mod tests {
assert!(listing.next().await.is_none());
}
#[tokio::test]
async fn next_page_preserves_empty_pages_and_remaining_items() {
let requests = Arc::new(Mutex::new(0));
let count = requests.clone();
let mut listing = Listing::new(ListingOptions::default(), move |options| {
*count.lock().unwrap() += 1;
async move {
let (items, page_token) = match options.page_token.as_deref() {
None => (vec![], Some("second".to_string())),
Some("second") => (vec!["a", "b"], Some("last".to_string())),
Some("last") => (vec!["c"], None),
_ => unreachable!(),
};
Ok(ListingPage { items, page_token })
}
});
assert_eq!(*requests.lock().unwrap(), 0);
assert!(listing.next_page().await.unwrap().unwrap().items.is_empty());
assert_eq!(*requests.lock().unwrap(), 1);
assert_eq!(listing.try_next().await.unwrap(), Some("a"));
assert_eq!(*requests.lock().unwrap(), 2);
let remaining = listing.next_page().await.unwrap().unwrap();
assert_eq!(remaining.items, vec!["b"]);
assert_eq!(remaining.page_token.as_deref(), Some("last"));
assert_eq!(*requests.lock().unwrap(), 2);
assert_eq!(listing.next_page().await.unwrap().unwrap().items, vec!["c"]);
assert!(listing.next_page().await.unwrap().is_none());
assert_eq!(*requests.lock().unwrap(), 3);
}
#[tokio::test]
async fn invalid_limits_do_not_fetch() {
for limit in [0, u32::MAX] {
rust/lancedb/src/remote/db.rs
+20 -46
View File
@@ -2349,16 +2349,8 @@ mod tests {
.unwrap()
.is_empty()
);
let page = conn
.list_tables(lance_namespace::models::ListTablesRequest {
id: Some(namespace),
limit: Some(0),
..Default::default()
})
.await
.unwrap();
assert!(page.tables.is_empty());
assert!(page.page_token.is_none());
let mut listing = conn.list_tables(&namespace, ListingOptions::default().page_limit(0));
assert!(listing.next_page().await.is_err());
}
#[rstest::rstest]
@@ -2393,25 +2385,17 @@ mod tests {
)
.unwrap()
});
let first = conn
.list_tables(lance_namespace::models::ListTablesRequest {
id: Some(namespace.clone()),
..Default::default()
})
.await
.unwrap();
assert_eq!(first.tables.len(), 100);
let mut first_listing = conn.list_tables(&namespace, ListingOptions::default());
let first = first_listing.next_page().await.unwrap().unwrap();
assert_eq!(first.items.len(), 100);
assert_eq!(first.page_token.as_deref(), Some("opaque-token"));
let second = conn
.list_tables(lance_namespace::models::ListTablesRequest {
id: Some(namespace),
page_token: first.page_token,
..Default::default()
})
.await
.unwrap();
let mut second_listing = conn.list_tables(
&namespace,
ListingOptions::default().page_token(first.page_token.unwrap()),
);
let second = second_listing.next_page().await.unwrap().unwrap();
let expected: Vec<_> = (100..130).map(|i| format!("t{i:03}")).collect();
assert_eq!(second.tables, expected);
assert_eq!(second.items, expected);
assert!(second.page_token.is_none());
}
@@ -3680,14 +3664,9 @@ mod tests {
assert!(table.is_ok(), "Failed to create table: {:?}", table.err());
// List tables in the child namespace
let list_response = conn
.list_tables(ListTablesRequest {
id: Some(namespace.clone()),
..Default::default()
})
.await
.expect("Failed to list tables");
assert_eq!(list_response.tables, vec!["test_table"]);
let list_response = conn.list_tables(&namespace, ListingOptions::default());
let names = list_response.try_collect::<Vec<_>>().await.unwrap();
assert_eq!(names, vec!["test_table"]);
// Get namespace client and verify it can also list tables
let namespace_client = conn.namespace_client().await.unwrap();
@@ -3754,17 +3733,12 @@ mod tests {
}
// List tables in the child namespace
let list_response = conn
.list_tables(ListTablesRequest {
id: Some(namespace.clone()),
..Default::default()
})
.await
.unwrap();
assert_eq!(list_response.tables.len(), 3);
assert!(list_response.tables.contains(&"table1".to_string()));
assert!(list_response.tables.contains(&"table2".to_string()));
assert!(list_response.tables.contains(&"table3".to_string()));
let list_response = conn.list_tables(&namespace, ListingOptions::default());
let names = list_response.try_collect::<Vec<_>>().await.unwrap();
assert_eq!(names.len(), 3);
assert!(names.contains(&"table1".to_string()));
assert!(names.contains(&"table2".to_string()));
assert!(names.contains(&"table3".to_string()));
}
}