Compare commits

...

3 Commits

Author SHA1 Message Date
Gatefixer e0d6b9a4fa fix: preserve seeded IVF-PQ training contracts 2026-08-06 07:35:47 +00:00
Gatefixer 845868a343 fix: harden deterministic IVF-PQ training 2026-08-06 06:32:58 +00:00
Gatefixer cc321e9801 fix(python): support deterministic IVF-PQ training 2026-08-06 05:00:20 +00:00
10 changed files with 1471 additions and 13 deletions
Generated
+1
View File
@@ -5442,6 +5442,7 @@ dependencies = [
"pprof 0.14.1",
"rand 0.9.5",
"random_word",
"rayon",
"regex",
"reqwest 0.12.28",
"rstest",
+1
View File
@@ -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"
+14
View File
@@ -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:
+5
View File
@@ -2762,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",
+12 -2
View File
@@ -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))
+4
View File
@@ -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>,
}
+1
View File
@@ -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" }
+27
View File
@@ -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 {
+41
View File
@@ -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