diff --git a/python/Cargo.toml b/python/Cargo.toml index 6181f704f..ddc13b69f 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -26,7 +26,9 @@ lance-namespace-impls.workspace = true lance-io.workspace = true env_logger.workspace = true log.workspace = true -pyo3 = { version = "0.28", features = ["extension-module", "abi3-py310", "chrono"] } +# Maturin enables extension-module mode for Python builds. Keeping it out of +# Cargo features lets Rust unit tests link against libpython. +pyo3 = { version = "0.28", features = ["abi3-py310", "chrono"] } chrono.workspace = true pyo3-async-runtimes = { version = "0.28", features = [ "attributes", @@ -41,10 +43,7 @@ tokio.workspace = true libc = "0.2" [build-dependencies] -pyo3-build-config = { version = "0.28", features = [ - "extension-module", - "abi3-py310", -] } +pyo3-build-config = { version = "0.28", features = ["abi3-py310"] } [features] default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/cos", "lancedb/goosefs", "lancedb/metrics-otel"] diff --git a/python/pyproject.toml b/python/pyproject.toml index ae42172c0..fad3b1001 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -103,7 +103,7 @@ python-source = "python" module-name = "lancedb._lancedb" [build-system] -requires = ["maturin>=1.4"] +requires = ["maturin>=1.9.4"] build-backend = "maturin" [tool.ruff.lint] diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index aa7846892..d2b63baf6 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -163,6 +163,15 @@ class FTS: The number of documents per compressed posting block. Supported values are 128 and 256. A value of 256 uses the experimental FTS V3 format and may introduce breaking changes. + memory_limit : int, optional + The total memory limit in MiB for the local FTS build stage. The limit + is divided evenly among indexing workers. This build-only setting is + not persisted with the index and does not apply to remote tables. + num_workers : int, optional + The number of workers for a local FTS build. By default Lance uses + roughly half of the available CPU cores. The effective value is + limited by the available compute capacity. This build-only setting is + not persisted with the index and does not apply to remote tables. Notes ----- @@ -185,6 +194,8 @@ class FTS: prefix_only: bool = False block_size: int = 128 custom_stop_words: Optional[List[str]] = None + memory_limit: Optional[int] = None + num_workers: Optional[int] = None @dataclass diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index f791f9886..625198d92 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -245,6 +245,14 @@ def test_create_inverted_index_rejects_invalid_block_size(table): table.create_index("text", config=FTS(block_size=129)) +def test_create_inverted_index_respects_build_memory_limit(table): + with pytest.raises(ValueError, match="exceeds worker memory limit"): + table.create_index( + "text", + config=FTS(memory_limit=0, num_workers=1), + ) + + def test_custom_stop_words_list(table): table.create_index( "text", diff --git a/python/src/index.rs b/python/src/index.rs index dd362373e..a5ca63c68 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -42,7 +42,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult Ok(LanceDbIndex::Fm(FmIndexBuilder::default())), "FTS" => { let params = source.extract::()?; - let inner_opts = FtsIndexBuilder::default() + let mut inner_opts = FtsIndexBuilder::default() .base_tokenizer(params.base_tokenizer) .language(¶ms.language) .map_err(|_| { @@ -61,6 +61,12 @@ pub fn extract_index_params(source: &Option>) -> PyResult, + num_workers: Option, } #[derive(FromPyObject)] @@ -444,3 +452,51 @@ impl IndexConfig { } } } + +#[cfg(test)] +mod tests { + use super::*; + use pyo3::types::{PyDict, PyDictMethods}; + use serde_json::json; + + #[test] + fn fts_build_controls_are_forwarded() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c"class FTS: + with_position = True + base_tokenizer = 'simple' + language = 'English' + max_token_length = None + lower_case = True + stem = False + remove_stop_words = False + custom_stop_words = None + ascii_folding = False + ngram_min_length = 3 + ngram_max_length = 3 + prefix_only = False + block_size = 128 + memory_limit = 2048 + num_workers = 7 + +config = FTS()", + None, + Some(&locals), + ) + .unwrap(); + + let config = locals.get_item("config").unwrap().unwrap(); + let index = extract_index_params(&Some(config)).unwrap(); + let LanceDbIndex::FTS(params) = index else { + panic!("expected FTS index parameters"); + }; + let training_json = params.to_training_json().unwrap(); + + assert_eq!(training_json.get("memory_limit"), Some(&json!(2048))); + assert_eq!(training_json.get("num_workers"), Some(&json!(7))); + }); + } +}