mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-31 10:38:31 +00:00
Compare commits
5 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| e0d6b9a4fa | |||
| 845868a343 | |||
| cc321e9801 | |||
| 7357d63e87 | |||
| 624a75edf7 |
Generated
+1
@@ -5442,6 +5442,7 @@ dependencies = [
|
||||
"pprof 0.14.1",
|
||||
"rand 0.9.5",
|
||||
"random_word",
|
||||
"rayon",
|
||||
"regex",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
|
||||
@@ -60,6 +60,7 @@ moka = { version = "0.12", features = ["future"] }
|
||||
object_store = "0.13.2"
|
||||
pin-project = "1.0.7"
|
||||
rand = "0.9"
|
||||
rayon = "1"
|
||||
snafu = "0.8"
|
||||
url = "2"
|
||||
num-traits = "0.2"
|
||||
|
||||
@@ -707,6 +707,9 @@ class LanceDBConnection(DBConnection):
|
||||
self._namespace_client_properties = namespace_client_properties
|
||||
if _inner is not None:
|
||||
self._conn = _inner
|
||||
# Native-derived wrappers resolve this in their async reconstruction
|
||||
# path so construction never synchronously re-enters LOOP.
|
||||
self._read_consistency_interval = read_consistency_interval
|
||||
self._cached_namespace_client = None
|
||||
return
|
||||
|
||||
@@ -756,11 +759,14 @@ class LanceDBConnection(DBConnection):
|
||||
# storage_options. Also, this class really shouldn't be holding any state
|
||||
# beyond _conn.
|
||||
self._conn = AsyncConnection(LOOP.run(do_connect()))
|
||||
# Keep property access synchronous so debugger introspection cannot wait on
|
||||
# the background loop while that thread is suspended at a breakpoint.
|
||||
self._read_consistency_interval = read_consistency_interval
|
||||
self._cached_namespace_client: Optional[LanceNamespace] = None
|
||||
|
||||
@property
|
||||
def read_consistency_interval(self) -> Optional[timedelta]:
|
||||
return LOOP.run(self._conn.get_read_consistency_interval())
|
||||
return self._read_consistency_interval
|
||||
|
||||
@property
|
||||
def session(self) -> Optional[Session]:
|
||||
@@ -771,8 +777,16 @@ class LanceDBConnection(DBConnection):
|
||||
return self._conn.uri
|
||||
|
||||
@classmethod
|
||||
def from_inner(cls, inner: LanceDbConnection):
|
||||
return cls(None, _inner=inner)
|
||||
def from_inner(
|
||||
cls,
|
||||
inner: LanceDbConnection,
|
||||
read_consistency_interval: Optional[timedelta],
|
||||
):
|
||||
return cls(
|
||||
None,
|
||||
read_consistency_interval=read_consistency_interval,
|
||||
_inner=inner,
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(uri={self._conn.uri!r})"
|
||||
|
||||
@@ -769,6 +769,13 @@ class IvfPq:
|
||||
|
||||
The default value is 256.
|
||||
|
||||
seed: int, optional
|
||||
Seed used for deterministic sampling and training. Given identical data in
|
||||
the same row order and identical index parameters, the same seed produces
|
||||
the same IVF centroids and PQ codebook. This option is supported for local
|
||||
CPU index builds; remote and accelerator builds reject it explicitly. If
|
||||
omitted, training remains random.
|
||||
|
||||
target_partition_size: int, default is 8192
|
||||
|
||||
The target size of each partition.
|
||||
@@ -783,11 +790,18 @@ class IvfPq:
|
||||
num_bits: int = 8
|
||||
max_iterations: int = 50
|
||||
sample_rate: int = 256
|
||||
seed: Optional[int] = None
|
||||
target_partition_size: Optional[int] = None
|
||||
# Name of the accelerator (e.g. "cuda") to use for IVF training. When set,
|
||||
# create_index() dispatches to pylance to build the index on the accelerator.
|
||||
accelerator: Optional[str] = None
|
||||
|
||||
def __post_init__(self):
|
||||
if self.seed is not None and self.accelerator is not None:
|
||||
raise ValueError(
|
||||
"IvfPq seed is not supported with accelerator-based index training"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class IvfRq:
|
||||
|
||||
@@ -226,7 +226,7 @@ class PermutationBuilder:
|
||||
|
||||
async def do_execute():
|
||||
inner_tbl = await self._async.execute()
|
||||
return LanceTable.from_inner(inner_tbl)
|
||||
return await LanceTable.from_inner(inner_tbl)
|
||||
|
||||
return LOOP.run(do_execute())
|
||||
|
||||
|
||||
@@ -2182,11 +2182,15 @@ class LanceTable(Table):
|
||||
return self.name
|
||||
|
||||
@classmethod
|
||||
def from_inner(cls, tbl: LanceDBTable):
|
||||
from .db import LanceDBConnection
|
||||
async def from_inner(cls, tbl: LanceDBTable):
|
||||
from .db import AsyncConnection, LanceDBConnection
|
||||
|
||||
async_tbl = AsyncTable(tbl)
|
||||
conn = LanceDBConnection.from_inner(tbl.database())
|
||||
inner_conn = tbl.database()
|
||||
read_consistency_interval = await AsyncConnection(
|
||||
inner_conn
|
||||
).get_read_consistency_interval()
|
||||
conn = LanceDBConnection.from_inner(inner_conn, read_consistency_interval)
|
||||
return cls(
|
||||
conn,
|
||||
async_tbl.name,
|
||||
@@ -2758,6 +2762,11 @@ class LanceTable(Table):
|
||||
if config is not None and hasattr(config, "accelerator"):
|
||||
acc = getattr(config, "accelerator", None)
|
||||
if acc is not None:
|
||||
if isinstance(config, IvfPq) and config.seed is not None:
|
||||
raise ValueError(
|
||||
"IvfPq seed is not supported with accelerator-based "
|
||||
"index training"
|
||||
)
|
||||
# Dispatch to pylance for GPU acceleration
|
||||
index_type_map = {
|
||||
"IvfFlat": "IVF_FLAT",
|
||||
|
||||
@@ -77,6 +77,23 @@ def test_sync_repr_does_not_use_background_loop(tmp_path, monkeypatch):
|
||||
assert repr(table) == f"LanceTable(name='test', _conn={db!r})"
|
||||
|
||||
|
||||
def test_read_consistency_interval_does_not_use_background_loop(tmp_path, monkeypatch):
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.db import LanceDBConnection
|
||||
|
||||
consistency_interval = timedelta(seconds=5)
|
||||
db = lancedb.connect(tmp_path, read_consistency_interval=consistency_interval)
|
||||
db_from_inner = LanceDBConnection.from_inner(db._inner, consistency_interval)
|
||||
|
||||
def fail_run(*args, **kwargs):
|
||||
raise AssertionError("properties should not use the Python background loop")
|
||||
|
||||
monkeypatch.setattr(LOOP, "run", fail_run)
|
||||
|
||||
assert db.read_consistency_interval == consistency_interval
|
||||
assert db_from_inner.read_consistency_interval == consistency_interval
|
||||
|
||||
|
||||
def test_ingest_pd(tmp_path):
|
||||
db = lancedb.connect(tmp_path)
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@ from lancedb.index import (
|
||||
HnswFlat,
|
||||
FTS,
|
||||
)
|
||||
from lancedb.table import IndexStatistics
|
||||
from lancedb.table import IndexStatistics, LanceTable
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
@@ -375,7 +375,7 @@ async def test_create_vector_index(some_table: AsyncTable):
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
|
||||
# Can create
|
||||
await some_table.create_index("vector", config=IvfPq(num_bits=4))
|
||||
await some_table.create_index("vector", config=IvfPq(num_bits=4, seed=42))
|
||||
# Can recreate if replace=True
|
||||
await some_table.create_index("vector", config=IvfPq(num_bits=4), replace=True)
|
||||
# Can't recreate if replace=False
|
||||
@@ -395,6 +395,16 @@ async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
|
||||
assert stats.num_indices == 1
|
||||
|
||||
|
||||
def test_seeded_ivfpq_rejects_accelerator():
|
||||
with pytest.raises(ValueError, match="seed is not supported with accelerator"):
|
||||
IvfPq(seed=42, accelerator="cuda")
|
||||
|
||||
config = IvfPq(seed=42)
|
||||
config.accelerator = "cuda"
|
||||
with pytest.raises(ValueError, match="seed is not supported with accelerator"):
|
||||
object.__new__(LanceTable).create_index("vector", config=config)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_ivfrq_index(some_table: AsyncTable):
|
||||
await some_table.create_index("vector", config=IvfRq(num_bits=1))
|
||||
|
||||
@@ -6,6 +6,7 @@ import math
|
||||
import pytest
|
||||
|
||||
from lancedb import DBConnection, Table, connect
|
||||
from lancedb.background_loop import LOOP
|
||||
from lancedb.permutation import Permutation, Permutations, permutation_builder
|
||||
|
||||
|
||||
@@ -31,6 +32,25 @@ def test_split_random_ratios(mem_db):
|
||||
assert 65 <= split_1_count <= 75 # ~70% ± tolerance
|
||||
|
||||
|
||||
def test_execute_does_not_reenter_background_loop(tmp_path, monkeypatch):
|
||||
import threading
|
||||
|
||||
db = connect(tmp_path)
|
||||
tbl = db.create_table("test_table", pa.table({"x": range(10)}))
|
||||
original_run = LOOP.run
|
||||
|
||||
def fail_on_reentry(future):
|
||||
assert threading.current_thread() is not LOOP.thread
|
||||
return original_run(future)
|
||||
|
||||
monkeypatch.setattr(LOOP, "run", fail_on_reentry)
|
||||
|
||||
permutation_tbl = permutation_builder(tbl).execute()
|
||||
|
||||
assert permutation_tbl.count_rows() == 10
|
||||
assert permutation_tbl._conn.read_consistency_interval is None
|
||||
|
||||
|
||||
def test_split_random_counts(mem_db):
|
||||
"""Test random splitting with absolute counts."""
|
||||
tbl = mem_db.create_table(
|
||||
|
||||
@@ -6,6 +6,7 @@ import os
|
||||
import sys
|
||||
import threading
|
||||
import warnings
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import date, datetime, timedelta
|
||||
from time import sleep
|
||||
from typing import List
|
||||
@@ -2124,6 +2125,27 @@ def test_delete(mem_db: DBConnection):
|
||||
assert table.to_arrow()["id"].to_pylist() == [1]
|
||||
|
||||
|
||||
def test_concurrent_deletes_are_thread_safe(mem_db: DBConnection):
|
||||
num_workers = 8
|
||||
table = mem_db.create_table(
|
||||
"my_table", data=[{"id": row_id} for row_id in range(num_workers)]
|
||||
)
|
||||
barrier = threading.Barrier(num_workers)
|
||||
|
||||
def delete(row_id: int):
|
||||
barrier.wait()
|
||||
return table.delete(f"id = {row_id}")
|
||||
|
||||
with ThreadPoolExecutor(max_workers=num_workers) as pool:
|
||||
results = list(pool.map(delete, range(num_workers)))
|
||||
|
||||
assert all(result.num_deleted_rows == 1 for result in results)
|
||||
assert sorted(result.version for result in results) == list(
|
||||
range(2, num_workers + 2)
|
||||
)
|
||||
assert table.count_rows() == 0
|
||||
|
||||
|
||||
def test_delete_expr(mem_db: DBConnection):
|
||||
table = mem_db.create_table(
|
||||
"my_table",
|
||||
|
||||
@@ -90,6 +90,9 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
|
||||
.max_iterations(params.max_iterations)
|
||||
.sample_rate(params.sample_rate)
|
||||
.num_bits(params.num_bits);
|
||||
if let Some(seed) = params.seed {
|
||||
ivf_pq_builder = ivf_pq_builder.seed(seed);
|
||||
}
|
||||
if let Some(num_partitions) = params.num_partitions {
|
||||
ivf_pq_builder = ivf_pq_builder.num_partitions(num_partitions);
|
||||
}
|
||||
@@ -232,6 +235,7 @@ struct IvfPqParams {
|
||||
num_bits: u32,
|
||||
max_iterations: u32,
|
||||
sample_rate: u32,
|
||||
seed: Option<u64>,
|
||||
target_partition_size: Option<u32>,
|
||||
}
|
||||
|
||||
|
||||
@@ -745,6 +745,9 @@ impl Table {
|
||||
|
||||
#[allow(private_interfaces)]
|
||||
pub fn delete(self_: PyRef<'_, Self>, condition: PredicateArg) -> PyResult<Bound<'_, PyAny>> {
|
||||
// Do not hold the Python borrow across the await. The cloned Rust table
|
||||
// handle is thread-safe and allows deletes on the same Python table to
|
||||
// run concurrently without PyO3 reporting "Already borrowed".
|
||||
let inner = self_.inner_ref()?.clone();
|
||||
future_into_py(self_.py(), async move {
|
||||
let result = match &condition {
|
||||
|
||||
@@ -61,6 +61,7 @@ futures.workspace = true
|
||||
num-traits.workspace = true
|
||||
url.workspace = true
|
||||
rand.workspace = true
|
||||
rayon.workspace = true
|
||||
regex.workspace = true
|
||||
serde = { version = "^1" }
|
||||
serde_json = { version = "1" }
|
||||
|
||||
@@ -274,6 +274,8 @@ pub struct IvfPqIndexBuilder {
|
||||
pub(crate) sample_rate: u32,
|
||||
pub(crate) max_iterations: u32,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) seed: Option<u64>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) target_partition_size: Option<u32>,
|
||||
|
||||
// PQ
|
||||
@@ -292,6 +294,7 @@ impl Default for IvfPqIndexBuilder {
|
||||
num_bits: None,
|
||||
sample_rate: 256,
|
||||
max_iterations: 50,
|
||||
seed: None,
|
||||
target_partition_size: None,
|
||||
}
|
||||
}
|
||||
@@ -301,6 +304,30 @@ impl IvfPqIndexBuilder {
|
||||
impl_distance_type_setter!();
|
||||
impl_ivf_params_setter!();
|
||||
impl_pq_params_setter!();
|
||||
|
||||
/// Use a deterministic seed when sampling and training the IVF and PQ models.
|
||||
///
|
||||
/// Given identical data in the same row order and identical index parameters,
|
||||
/// using the same seed produces the same IVF centroids and PQ codebook. This is
|
||||
/// useful when independently-built tables need reproducible approximate-search
|
||||
/// results.
|
||||
///
|
||||
/// Seeded training is supported by native tables. Remote backends reject this
|
||||
/// option unless they can provide the same deterministic training contract.
|
||||
///
|
||||
/// If no seed is provided, index training uses random sampling and initialization.
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
/// use lancedb::index::vector::IvfPqIndexBuilder;
|
||||
///
|
||||
/// let index = IvfPqIndexBuilder::default().seed(42);
|
||||
/// ```
|
||||
pub fn seed(mut self, seed: u64) -> Self {
|
||||
self.seed = Some(seed);
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn suggested_num_sub_vectors(dim: u32) -> u32 {
|
||||
|
||||
@@ -339,6 +339,12 @@ impl<S: HttpSend> RemoteTable<S> {
|
||||
// Auto is special-cased since it needs schema inspection.
|
||||
let (index_type_str, params) = match &index.index {
|
||||
Index::IvfFlat(p) => ("IVF_FLAT", Some(to_json(p)?)),
|
||||
Index::IvfPq(p) if p.seed.is_some() => {
|
||||
return Err(Error::NotSupported {
|
||||
message: "Deterministic IVF PQ training is not supported on remote tables"
|
||||
.to_string(),
|
||||
});
|
||||
}
|
||||
Index::IvfPq(p) => ("IVF_PQ", Some(to_json(p)?)),
|
||||
Index::IvfSq(p) => ("IVF_SQ", Some(to_json(p)?)),
|
||||
Index::IvfHnswSq(p) => ("IVF_HNSW_SQ", Some(to_json(p)?)),
|
||||
@@ -5157,6 +5163,41 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_seeded_ivf_pq_is_rejected_for_remote_tables() {
|
||||
let table = Table::new_with_handler("my_table", |request| match request.url().path() {
|
||||
"/v1/table/my_table/describe/" => {
|
||||
let schema = Schema::new(vec![Field::new(
|
||||
"vector",
|
||||
DataType::FixedSizeList(
|
||||
Arc::new(Field::new("item", DataType::Float32, true)),
|
||||
8,
|
||||
),
|
||||
false,
|
||||
)]);
|
||||
http::Response::builder()
|
||||
.status(200)
|
||||
.body(describe_response(&schema))
|
||||
.unwrap()
|
||||
}
|
||||
path => panic!("Unexpected request for unsupported seeded index: {path}"),
|
||||
});
|
||||
|
||||
let error = table
|
||||
.create_index(
|
||||
&["vector"],
|
||||
Index::IvfPq(IvfPqIndexBuilder::default().seed(42)),
|
||||
)
|
||||
.execute()
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::NotSupported { message }
|
||||
if message.contains("not supported on remote tables")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_create_index_returns_job() {
|
||||
let describe_calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user