diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index a93c32e63..e0f6635d8 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -5,7 +5,7 @@ import os from concurrent.futures import ThreadPoolExecutor from datetime import timedelta -from typing import Dict, Optional, Union, Any, List, Iterable, TYPE_CHECKING +from typing import Dict, Optional, Union, Any, List, Iterable, Literal, TYPE_CHECKING from packaging.version import Version as _Version @@ -354,13 +354,31 @@ def tokenize( ngram_min_length: int = 3, ngram_max_length: int = 3, prefix_only: bool = False, + analyzer: Optional[Literal["text", "code"]] = None, + split_identifiers: Optional[bool] = None, + split_on_numerics: Optional[bool] = None, + preserve_original: Optional[bool] = None, + index_operators: Optional[bool] = None, ) -> Iterable[FtsToken]: """Tokenize a full-text search query using an explicit tokenizer. This does not require an FTS index. The tokenizer options match :class:`lancedb.index.FTS`. ``custom_stop_words`` accepts a list of strings. When ``stem`` or ``remove_stop_words`` is None, the base tokenizer's default - applies: False for ``code`` and ``ngram``, True otherwise. + applies: False for ``code`` and ``ngram``, True otherwise. ``analyzer`` and the + code flags (``split_identifiers``, ``split_on_numerics``, ``preserve_original`` + and ``index_operators``) behave as they do on :class:`lancedb.index.FTS`. + + Examples + -------- + Inspect the identifier and subwords produced by the code analyzer: + + >>> import lancedb + >>> tokens = lancedb.tokenize( + ... "parseHTTPRequest", analyzer="code", split_identifiers=True + ... ) + >>> [token.text for token in tokens] + ['parsehttprequest', 'parse', 'http', 'request'] """ return _tokenize( @@ -376,6 +394,11 @@ def tokenize( ngram_min_length=ngram_min_length, ngram_max_length=ngram_max_length, prefix_only=prefix_only, + analyzer=analyzer, + split_identifiers=split_identifiers, + split_on_numerics=split_on_numerics, + preserve_original=preserve_original, + index_operators=index_operators, ) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 8730917d3..33689f900 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -67,6 +67,11 @@ def tokenize( ngram_min_length: int = 3, ngram_max_length: int = 3, prefix_only: bool = False, + analyzer: Optional[Literal["text", "code"]] = None, + split_identifiers: Optional[bool] = None, + split_on_numerics: Optional[bool] = None, + preserve_original: Optional[bool] = None, + index_operators: Optional[bool] = None, ) -> List["FtsToken"]: ... class PyExpr: diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index d4c7f8c8c..f0700d0b2 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -182,6 +182,10 @@ class FTS: >>> from lancedb.query import DocumentGranularity >>> config = FTS(document_granularity=DocumentGranularity.LIST_ELEMENT) + Create an index that splits code identifiers into searchable subwords: + + >>> config = FTS(analyzer="code", split_identifiers=True) + Attributes ---------- with_position : bool, default False @@ -243,6 +247,28 @@ class FTS: ``LIST_ELEMENT`` treats each element of the deepest list on the indexed field path as one document and returns its physical coordinates in ``_doc_index`` for matching queries. + analyzer : {"text", "code"}, optional + ``"code"`` selects the code analyzer. ``"text"`` selects a text-family + analyzer with the selected ``base_tokenizer``; it does not reset an + explicit compatible tokenizer to ``"simple"``. ``None`` infers the + analyzer from ``base_tokenizer``. ``"code"`` cannot be combined with a + ``base_tokenizer`` other than ``"simple"`` (the default) or ``"code"``. + ``"text"`` cannot be combined with ``"code"``. + split_identifiers : bool, optional + Split code identifiers such as ``camelCase`` and ``snake_case`` into + subwords. Setting this to True requires the code analyzer. The default + is False, so ``analyzer="code"`` alone keeps identifiers whole. + split_on_numerics : bool, optional + Split identifier subwords at letter and number boundaries, for example + ``utf8`` into ``utf`` and ``8``. The code analyzer default is True, but + this only takes effect when ``split_identifiers=True``. + preserve_original : bool, optional + Also index the complete identifier next to its subwords. The code + analyzer default is True, but this only takes effect when + ``split_identifiers=True``. + index_operators : bool, optional + Index code operators such as ``::``, ``->`` and ``!=``. The code + analyzer default is False. Notes ----- @@ -268,6 +294,11 @@ class FTS: memory_limit: Optional[int] = None num_workers: Optional[int] = None document_granularity: DocumentGranularity = DocumentGranularity.ROW + analyzer: Optional[Literal["text", "code"]] = None + split_identifiers: Optional[bool] = None + split_on_numerics: Optional[bool] = None + preserve_original: Optional[bool] = None + index_operators: Optional[bool] = None @dataclass diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index 6714f35d9..a6e55a4c3 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -1832,3 +1832,120 @@ async def test_fts_fast_search_async(async_table): .to_list() ) assert len(results) > 0 + + +CODE_TEXT = "parseHTTPRequest user_id utf8" + + +def _token_texts(tokens): + return [token.text for token in tokens] + + +def test_tokenize_code_analyzer_flags(): + default = _token_texts(ldb.tokenize(CODE_TEXT, analyzer="code")) + assert default == _token_texts(ldb.tokenize(CODE_TEXT, base_tokenizer="code")) + + no_split = _token_texts( + ldb.tokenize(CODE_TEXT, analyzer="code", split_identifiers=False) + ) + split = _token_texts( + ldb.tokenize(CODE_TEXT, analyzer="code", split_identifiers=True) + ) + assert "parsehttprequest" in no_split + assert "parse" not in no_split + assert {"parse", "http", "request", "user", "id"} <= set(split) + + numerics = _token_texts( + ldb.tokenize( + CODE_TEXT, + analyzer="code", + split_identifiers=True, + split_on_numerics=True, + ) + ) + no_numerics = _token_texts( + ldb.tokenize( + CODE_TEXT, + analyzer="code", + split_identifiers=True, + split_on_numerics=False, + ) + ) + assert "utf" in numerics and "8" in numerics + assert "utf8" in no_numerics + + operators = _token_texts( + ldb.tokenize("a::b -> c != d", analyzer="code", index_operators=True) + ) + plain = _token_texts( + ldb.tokenize("a::b -> c != d", analyzer="code", index_operators=False) + ) + assert {"::", "->", "!="} <= set(operators) + assert not {"::", "->", "!="} & set(plain) + + +@pytest.mark.parametrize( + ("kwargs", "message"), + [ + ({"analyzer": "code", "base_tokenizer": "ngram"}, "requires base_tokenizer"), + ({"analyzer": "text", "base_tokenizer": "code"}, "requires analyzer='code'"), + ({"analyzer": "java"}, "analyzer must be"), + ({"split_identifiers": True}, "require analyzer='code'"), + ({"base_tokenizer": "raw", "index_operators": True}, "require analyzer='code'"), + ], +) +def test_code_analyzer_options_are_validated(mem_db, kwargs, message): + with pytest.raises(ValueError, match=message): + ldb.tokenize(CODE_TEXT, **kwargs) + + table = mem_db.create_table("t", data=[{"text": CODE_TEXT}]) + with pytest.raises(ValueError, match=message): + table.create_index("text", config=FTS(**kwargs)) + + +def test_code_analyzer_flags_false_are_allowed_for_text(mem_db): + tokens = ldb.tokenize("Running in cafés", split_identifiers=False) + assert _token_texts(tokens) == ["run", "cafe"] + + +def test_fts_code_analyzer_index_matches_subwords(mem_db): + table = mem_db.create_table( + "code", + data=[{"text": "def parseHTTPRequest(user_id): pass"}, {"text": "unrelated"}], + ) + table.create_index( + "text", + config=FTS(analyzer="code", split_identifiers=True, preserve_original=True), + ) + + details = table.list_indices()[0].index_details + assert details["base_tokenizer"] == "code" + assert details["split_identifiers"] is True + assert details["preserve_original"] is True + + for query in ["http", "request", "parseHTTPRequest"]: + rows = table.search(query, query_type="fts").to_list() + assert [row["text"] for row in rows] == ["def parseHTTPRequest(user_id): pass"] + + +def test_fts_code_analyzer_without_splitting_does_not_match_subwords(mem_db): + table = mem_db.create_table( + "code", data=[{"text": "def parseHTTPRequest(user_id): pass"}] + ) + table.create_index("text", config=FTS(analyzer="code", split_identifiers=False)) + + assert table.search("request", query_type="fts").to_list() == [] + assert len(table.search("parseHTTPRequest", query_type="fts").to_list()) == 1 + + +@pytest.mark.asyncio +async def test_async_fts_code_analyzer(mem_db_async): + table = await mem_db_async.create_table( + "code", data=[{"text": "def parseHTTPRequest(user_id): pass"}] + ) + await table.create_index( + "text", config=FTS(analyzer="code", split_identifiers=True) + ) + + rows = await table.query().nearest_to_text("http").to_list() + assert len(rows) == 1 diff --git a/python/src/index.rs b/python/src/index.rs index d80f8d282..3ec03b6fd 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -35,6 +35,80 @@ pub fn class_name(ob: &'_ Bound<'_, PyAny>) -> PyResult { } } +/// Code-analyzer options shared by `FTS` and `tokenize`. +#[derive(Default)] +pub struct CodeAnalyzerOptions { + pub analyzer: Option, + pub split_identifiers: Option, + pub split_on_numerics: Option, + pub preserve_original: Option, + pub index_operators: Option, +} + +impl CodeAnalyzerOptions { + /// Check the options against the requested base tokenizer and return the + /// base tokenizer to use. `analyzer="code"` selects the code tokenizer when + /// `base_tokenizer` is left at its default of `"simple"`. + pub fn resolve_base_tokenizer(&self, base_tokenizer: String) -> PyResult { + let uses_code = match self.analyzer.as_deref() { + None | Some("text") => base_tokenizer == "code", + Some("code") => true, + Some(other) => { + return Err(PyValueError::new_err(format!( + "analyzer must be 'text' or 'code', got '{other}'" + ))); + } + }; + match self.analyzer.as_deref() { + Some("text") if base_tokenizer == "code" => { + return Err(PyValueError::new_err( + "base_tokenizer='code' requires analyzer='code'", + )); + } + Some("code") if base_tokenizer != "code" && base_tokenizer != "simple" => { + return Err(PyValueError::new_err(format!( + "analyzer='code' requires base_tokenizer='code', got '{base_tokenizer}'" + ))); + } + _ => {} + } + let flags = [ + self.split_identifiers, + self.split_on_numerics, + self.preserve_original, + self.index_operators, + ]; + if !uses_code && flags.into_iter().flatten().any(|value| value) { + return Err(PyValueError::new_err( + "split_identifiers, split_on_numerics, preserve_original and \ + index_operators require analyzer='code'", + )); + } + Ok(if uses_code { + "code".to_string() + } else { + base_tokenizer + }) + } + + /// Apply the explicit flags on top of the tokenizer defaults. + pub fn apply(&self, mut params: FtsIndexBuilder) -> FtsIndexBuilder { + if let Some(value) = self.split_identifiers { + params = params.split_identifiers(value); + } + if let Some(value) = self.split_on_numerics { + params = params.split_on_numerics(value); + } + if let Some(value) = self.preserve_original { + params = params.preserve_original(value); + } + if let Some(value) = self.index_operators { + params = params.index_operators(value); + } + params + } +} + pub fn extract_index_params(source: &Option>) -> PyResult { if let Some(source) = source { match class_name(source)?.as_str() { @@ -45,8 +119,16 @@ pub fn extract_index_params(source: &Option>) -> PyResult Ok(LanceDbIndex::NGram(NGramIndexBuilder::default())), "FTS" => { let params = source.extract::()?; + let code_options = CodeAnalyzerOptions { + analyzer: params.analyzer, + split_identifiers: params.split_identifiers, + split_on_numerics: params.split_on_numerics, + preserve_original: params.preserve_original, + index_operators: params.index_operators, + }; + let base_tokenizer = code_options.resolve_base_tokenizer(params.base_tokenizer)?; let mut inner_opts = FtsIndexBuilder::default() - .base_tokenizer(params.base_tokenizer) + .base_tokenizer(base_tokenizer) .language(¶ms.language) .map_err(|_| { PyValueError::new_err(format!( @@ -78,6 +160,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult, num_workers: Option, document_granularity: String, + analyzer: Option, + split_identifiers: Option, + split_on_numerics: Option, + preserve_original: Option, + index_operators: Option, } #[derive(FromPyObject)] @@ -496,6 +584,11 @@ mod tests { memory_limit = 2048 num_workers = 7 document_granularity = 'row' + analyzer = None + split_identifiers = None + split_on_numerics = None + preserve_original = None + index_operators = None config = FTS()", None, diff --git a/python/src/table.rs b/python/src/table.rs index 8d758697c..516e492b1 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -14,7 +14,7 @@ use crate::{ connection::Connection, error::PythonErrorExt, expr::PyExpr, - index::{IndexConfig, extract_index_params}, + index::{CodeAnalyzerOptions, IndexConfig, extract_index_params}, query::{Query, TakeQuery}, table::scannable::PyScannable, }; @@ -1030,7 +1030,12 @@ impl From for FtsToken { ascii_folding = true, ngram_min_length = 3, ngram_max_length = 3, - prefix_only = false + prefix_only = false, + analyzer = None, + split_identifiers = None, + split_on_numerics = None, + preserve_original = None, + index_operators = None ))] #[allow(clippy::too_many_arguments)] pub fn tokenize( @@ -1046,7 +1051,20 @@ pub fn tokenize( ngram_min_length: u32, ngram_max_length: u32, prefix_only: bool, + analyzer: Option, + split_identifiers: Option, + split_on_numerics: Option, + preserve_original: Option, + index_operators: Option, ) -> PyResult> { + let code_options = CodeAnalyzerOptions { + analyzer, + split_identifiers, + split_on_numerics, + preserve_original, + index_operators, + }; + let base_tokenizer = code_options.resolve_base_tokenizer(base_tokenizer)?; let mut params = FtsIndexBuilder::default() .base_tokenizer(base_tokenizer) .language(&language) @@ -1069,6 +1087,7 @@ pub fn tokenize( if let Some(remove_stop_words) = remove_stop_words { params = params.remove_stop_words(remove_stop_words); } + let params = code_options.apply(params); let tokens = lancedb_tokenize(&query, ¶ms).infer_error()?; Ok(tokens.into_iter().map(FtsToken::from).collect()) } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index e37e97755..2ff842319 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -8164,6 +8164,28 @@ mod tests { .document_granularity(DocumentGranularity::ListElement), ), ), + ( + "FTS", + { + let mut body = serde_json::to_value(InvertedIndexParams::default()).unwrap(); + body["base_tokenizer"] = "code".into(); + body["stem"] = false.into(); + body["remove_stop_words"] = false.into(); + body["split_identifiers"] = true.into(); + body["split_on_numerics"] = false.into(); + body["preserve_original"] = true.into(); + body["index_operators"] = true.into(); + body + }, + Index::FTS( + InvertedIndexParams::default() + .base_tokenizer("code".to_string()) + .split_identifiers(true) + .split_on_numerics(false) + .preserve_original(true) + .index_operators(true), + ), + ), ]; for (index_type, expected_body, index) in cases {