mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-13 14:16:54 +00:00
feat(python): expose code analyzer options on FTS and tokenize
Part of #4364 (Python side; the TypeScript `FtsOptions` change will follow in a separate PR). Lance already supports the code analyzer, but `FTS` and `lancedb.tokenize()` had no way to set it, so `FTS({'analyzer': 'code'})` and `FTS(split_identifiers=True)` raised `TypeError`. ## Changes - `FTS` and `lancedb.tokenize()` gain `analyzer` (`"text"` or `"code"`), `split_identifiers`, `split_on_numerics`, `preserve_original` and `index_operators`. They default to `None`, which keeps the tokenizer's own defaults, so existing indexes and calls are unchanged. - Validation follows pylance: `analyzer="code"` with a `base_tokenizer` other than `"code"` or the default `"simple"` is rejected, `analyzer="text"` with `base_tokenizer="code"` is rejected, and turning on a code flag without the code analyzer is rejected. Passing a flag as `False` is always allowed. - `analyzer="code"` with the default `base_tokenizer="simple"` selects the code tokenizer. `FTS` can't tell an explicit `"simple"` from the default, so that case is accepted too. - The deprecated `create_fts_index` is left alone; it points users at `create_index(config=FTS(...))`. - Rust `FtsIndexBuilder` is Lance's `InvertedIndexParams`, so the Rust SDK already had these options. The remote client serializes the same struct, so the options reach LanceDB Distributed in the `create_index` request body. A new case in the Rust `test_create_index` checks that body. ## Tests - `python/tests/test_fts.py`: tokenizer output for each flag, validation errors for `tokenize` and `FTS`, `False` flags on a text index, and local sync and async indexes that match `http`, `request` and `parseHTTPRequest` with splitting on and not match subwords with it off. - `cargo test --features remote remote::table::tests::test_create_index`
This commit is contained in:
+94
-1
@@ -35,6 +35,80 @@ pub fn class_name(ob: &'_ Bound<'_, PyAny>) -> PyResult<String> {
|
||||
}
|
||||
}
|
||||
|
||||
/// Code-analyzer options shared by `FTS` and `tokenize`.
|
||||
#[derive(Default)]
|
||||
pub struct CodeAnalyzerOptions {
|
||||
pub analyzer: Option<String>,
|
||||
pub split_identifiers: Option<bool>,
|
||||
pub split_on_numerics: Option<bool>,
|
||||
pub preserve_original: Option<bool>,
|
||||
pub index_operators: Option<bool>,
|
||||
}
|
||||
|
||||
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<String> {
|
||||
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<Bound<'_, PyAny>>) -> PyResult<LanceDbIndex> {
|
||||
if let Some(source) = source {
|
||||
match class_name(source)?.as_str() {
|
||||
@@ -45,8 +119,16 @@ pub fn extract_index_params(source: &Option<Bound<'_, PyAny>>) -> PyResult<Lance
|
||||
"NGram" => Ok(LanceDbIndex::NGram(NGramIndexBuilder::default())),
|
||||
"FTS" => {
|
||||
let params = source.extract::<FtsParams>()?;
|
||||
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<Bound<'_, PyAny>>) -> PyResult<Lance
|
||||
if let Some(num_workers) = params.num_workers {
|
||||
inner_opts = inner_opts.num_workers(num_workers);
|
||||
}
|
||||
let inner_opts = code_options.apply(inner_opts);
|
||||
let inner_opts = inner_opts
|
||||
.block_size(params.block_size)
|
||||
.map_err(|err| PyValueError::new_err(err.to_string()))?;
|
||||
@@ -235,6 +318,11 @@ struct FtsParams {
|
||||
memory_limit: Option<u64>,
|
||||
num_workers: Option<usize>,
|
||||
document_granularity: String,
|
||||
analyzer: Option<String>,
|
||||
split_identifiers: Option<bool>,
|
||||
split_on_numerics: Option<bool>,
|
||||
preserve_original: Option<bool>,
|
||||
index_operators: Option<bool>,
|
||||
}
|
||||
|
||||
#[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,
|
||||
|
||||
+21
-2
@@ -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<LanceDbFtsToken> 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<String>,
|
||||
split_identifiers: Option<bool>,
|
||||
split_on_numerics: Option<bool>,
|
||||
preserve_original: Option<bool>,
|
||||
index_operators: Option<bool>,
|
||||
) -> PyResult<Vec<FtsToken>> {
|
||||
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())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user