diff --git a/python/python/tests/test_s3.py b/python/python/tests/test_s3.py index 256ccb1d4..70b423adb 100644 --- a/python/python/tests/test_s3.py +++ b/python/python/tests/test_s3.py @@ -4,6 +4,7 @@ import asyncio import copy +from concurrent.futures import ThreadPoolExecutor from datetime import timedelta import threading @@ -86,6 +87,25 @@ def test_s3_lifecycle(s3_bucket: str): asyncio.run(test()) +@pytest.mark.s3_test +def test_concurrent_open_table(s3_bucket: str): + uri = f"s3://{s3_bucket}/test_concurrent_open_table" + db = lancedb.connect(uri, storage_options=copy.copy(CONFIG)) + db.create_table("test", pa.table({"x": [1, 2, 3]})) + + num_workers = 32 + barrier = threading.Barrier(num_workers) + + def open_and_count(_): + barrier.wait() + return db.open_table("test").count_rows() + + with ThreadPoolExecutor(max_workers=num_workers) as pool: + row_counts = list(pool.map(open_and_count, range(num_workers))) + + assert row_counts == [3] * num_workers + + @pytest.fixture() def kms_key(): kms = get_boto3_client("kms", endpoint_url=CONFIG["aws_endpoint"]) diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index 17dc82756..064b5d28f 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -1476,7 +1476,7 @@ mod tests { use crate::table::{AnyQuery, WriteOptions}; use arrow_array::{Int32Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; - use futures::{TryStreamExt, stream::once}; + use futures::{TryStreamExt, future::try_join_all, stream::once}; use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; @@ -1614,6 +1614,59 @@ mod tests { ); } + #[tokio::test] + async fn test_concurrent_open_table_reuses_connection_object_store() { + let tempdir = tempdir().unwrap(); + let uri = tempdir.path().to_str().unwrap(); + let session = Arc::new(lance::session::Session::default()); + let request = ConnectRequest { + uri: uri.to_string(), + #[cfg(feature = "remote")] + client_config: Default::default(), + options: Default::default(), + namespace_client_properties: Default::default(), + manifest_enabled: false, + read_consistency_interval: None, + session: Some(session.clone()), + }; + let db = ListingDatabase::connect_with_options(&request) + .await + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + db.create_table(CreateTableRequest { + name: "test".to_string(), + namespace_path: vec![], + data: Box::new(RecordBatch::new_empty(schema)) as Box, + mode: CreateTableMode::Create, + write_options: Default::default(), + location: None, + namespace_client: None, + }) + .await + .unwrap(); + + let before = session.store_registry().stats(); + let opened_tables = try_join_all((0..32).map(|_| { + db.open_table(OpenTableRequest { + name: "test".to_string(), + namespace_path: vec![], + index_cache_size: None, + lance_read_params: None, + location: None, + namespace_client: None, + managed_versioning: None, + }) + })) + .await + .unwrap(); + let after = session.store_registry().stats(); + + assert_eq!(opened_tables.len(), 32); + assert_eq!(after.misses, before.misses); + assert_eq!(after.active_stores, before.active_stores); + assert!(after.hits >= before.hits + 32); + } + #[tokio::test] async fn test_listing_database_root_ops_do_not_create_manifest() { let tempdir = tempdir().unwrap();