mirror of
https://github.com/lancedb/lancedb.git
synced 2026-09-12 00:02:21 +00:00
Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e0d6b9a4fa | ||
|
|
845868a343 | ||
|
|
cc321e9801 |
Generated
+1
@@ -5442,6 +5442,7 @@ dependencies = [
|
|||||||
"pprof 0.14.1",
|
"pprof 0.14.1",
|
||||||
"rand 0.9.5",
|
"rand 0.9.5",
|
||||||
"random_word",
|
"random_word",
|
||||||
|
"rayon",
|
||||||
"regex",
|
"regex",
|
||||||
"reqwest 0.12.28",
|
"reqwest 0.12.28",
|
||||||
"rstest",
|
"rstest",
|
||||||
|
|||||||
@@ -60,6 +60,7 @@ moka = { version = "0.12", features = ["future"] }
|
|||||||
object_store = "0.13.2"
|
object_store = "0.13.2"
|
||||||
pin-project = "1.0.7"
|
pin-project = "1.0.7"
|
||||||
rand = "0.9"
|
rand = "0.9"
|
||||||
|
rayon = "1"
|
||||||
snafu = "0.8"
|
snafu = "0.8"
|
||||||
url = "2"
|
url = "2"
|
||||||
num-traits = "0.2"
|
num-traits = "0.2"
|
||||||
|
|||||||
@@ -769,6 +769,13 @@ class IvfPq:
|
|||||||
|
|
||||||
The default value is 256.
|
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
|
target_partition_size: int, default is 8192
|
||||||
|
|
||||||
The target size of each partition.
|
The target size of each partition.
|
||||||
@@ -783,11 +790,18 @@ class IvfPq:
|
|||||||
num_bits: int = 8
|
num_bits: int = 8
|
||||||
max_iterations: int = 50
|
max_iterations: int = 50
|
||||||
sample_rate: int = 256
|
sample_rate: int = 256
|
||||||
|
seed: Optional[int] = None
|
||||||
target_partition_size: Optional[int] = None
|
target_partition_size: Optional[int] = None
|
||||||
# Name of the accelerator (e.g. "cuda") to use for IVF training. When set,
|
# 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.
|
# create_index() dispatches to pylance to build the index on the accelerator.
|
||||||
accelerator: Optional[str] = None
|
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
|
@dataclass
|
||||||
class IvfRq:
|
class IvfRq:
|
||||||
|
|||||||
@@ -2762,6 +2762,11 @@ class LanceTable(Table):
|
|||||||
if config is not None and hasattr(config, "accelerator"):
|
if config is not None and hasattr(config, "accelerator"):
|
||||||
acc = getattr(config, "accelerator", None)
|
acc = getattr(config, "accelerator", None)
|
||||||
if acc is not 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
|
# Dispatch to pylance for GPU acceleration
|
||||||
index_type_map = {
|
index_type_map = {
|
||||||
"IvfFlat": "IVF_FLAT",
|
"IvfFlat": "IVF_FLAT",
|
||||||
|
|||||||
@@ -26,7 +26,7 @@ from lancedb.index import (
|
|||||||
HnswFlat,
|
HnswFlat,
|
||||||
FTS,
|
FTS,
|
||||||
)
|
)
|
||||||
from lancedb.table import IndexStatistics
|
from lancedb.table import IndexStatistics, LanceTable
|
||||||
|
|
||||||
|
|
||||||
@pytest_asyncio.fixture
|
@pytest_asyncio.fixture
|
||||||
@@ -375,7 +375,7 @@ async def test_create_vector_index(some_table: AsyncTable):
|
|||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
|
async def test_create_4bit_ivfpq_index(some_table: AsyncTable):
|
||||||
# Can create
|
# 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
|
# Can recreate if replace=True
|
||||||
await some_table.create_index("vector", config=IvfPq(num_bits=4), replace=True)
|
await some_table.create_index("vector", config=IvfPq(num_bits=4), replace=True)
|
||||||
# Can't recreate if replace=False
|
# 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
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_create_ivfrq_index(some_table: AsyncTable):
|
async def test_create_ivfrq_index(some_table: AsyncTable):
|
||||||
await some_table.create_index("vector", config=IvfRq(num_bits=1))
|
await some_table.create_index("vector", config=IvfRq(num_bits=1))
|
||||||
|
|||||||
@@ -90,6 +90,9 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
|
|||||||
.max_iterations(params.max_iterations)
|
.max_iterations(params.max_iterations)
|
||||||
.sample_rate(params.sample_rate)
|
.sample_rate(params.sample_rate)
|
||||||
.num_bits(params.num_bits);
|
.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 {
|
if let Some(num_partitions) = params.num_partitions {
|
||||||
ivf_pq_builder = ivf_pq_builder.num_partitions(num_partitions);
|
ivf_pq_builder = ivf_pq_builder.num_partitions(num_partitions);
|
||||||
}
|
}
|
||||||
@@ -232,6 +235,7 @@ struct IvfPqParams {
|
|||||||
num_bits: u32,
|
num_bits: u32,
|
||||||
max_iterations: u32,
|
max_iterations: u32,
|
||||||
sample_rate: u32,
|
sample_rate: u32,
|
||||||
|
seed: Option<u64>,
|
||||||
target_partition_size: Option<u32>,
|
target_partition_size: Option<u32>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -61,6 +61,7 @@ futures.workspace = true
|
|||||||
num-traits.workspace = true
|
num-traits.workspace = true
|
||||||
url.workspace = true
|
url.workspace = true
|
||||||
rand.workspace = true
|
rand.workspace = true
|
||||||
|
rayon.workspace = true
|
||||||
regex.workspace = true
|
regex.workspace = true
|
||||||
serde = { version = "^1" }
|
serde = { version = "^1" }
|
||||||
serde_json = { version = "1" }
|
serde_json = { version = "1" }
|
||||||
|
|||||||
@@ -274,6 +274,8 @@ pub struct IvfPqIndexBuilder {
|
|||||||
pub(crate) sample_rate: u32,
|
pub(crate) sample_rate: u32,
|
||||||
pub(crate) max_iterations: u32,
|
pub(crate) max_iterations: u32,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[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>,
|
pub(crate) target_partition_size: Option<u32>,
|
||||||
|
|
||||||
// PQ
|
// PQ
|
||||||
@@ -292,6 +294,7 @@ impl Default for IvfPqIndexBuilder {
|
|||||||
num_bits: None,
|
num_bits: None,
|
||||||
sample_rate: 256,
|
sample_rate: 256,
|
||||||
max_iterations: 50,
|
max_iterations: 50,
|
||||||
|
seed: None,
|
||||||
target_partition_size: None,
|
target_partition_size: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -301,6 +304,30 @@ impl IvfPqIndexBuilder {
|
|||||||
impl_distance_type_setter!();
|
impl_distance_type_setter!();
|
||||||
impl_ivf_params_setter!();
|
impl_ivf_params_setter!();
|
||||||
impl_pq_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 {
|
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.
|
// Auto is special-cased since it needs schema inspection.
|
||||||
let (index_type_str, params) = match &index.index {
|
let (index_type_str, params) = match &index.index {
|
||||||
Index::IvfFlat(p) => ("IVF_FLAT", Some(to_json(p)?)),
|
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::IvfPq(p) => ("IVF_PQ", Some(to_json(p)?)),
|
||||||
Index::IvfSq(p) => ("IVF_SQ", Some(to_json(p)?)),
|
Index::IvfSq(p) => ("IVF_SQ", Some(to_json(p)?)),
|
||||||
Index::IvfHnswSq(p) => ("IVF_HNSW_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]
|
#[tokio::test]
|
||||||
async fn test_create_index_returns_job() {
|
async fn test_create_index_returns_job() {
|
||||||
let describe_calls = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
|
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