From 06b53c97d675a90498944036386b88cba9de3250 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Tue, 14 Jul 2026 10:59:33 -0700 Subject: [PATCH] feat: add table FTS query tokenization (#3659) ## Summary - add table-level FTS query tokenization returning token text and position - use the native index tokenizer for local tables and remote index metadata for remote tables - expose sync and async Python table wrappers with focused coverage --- docs/src/js/classes/Table.md | 26 +++ docs/src/js/functions/tokenize.md | 26 +++ docs/src/js/globals.md | 5 + docs/src/js/interfaces/FtsOptions.md | 6 +- docs/src/js/interfaces/FtsToken.md | 29 ++++ docs/src/js/interfaces/TokenizeOptions.md | 109 ++++++++++++ docs/src/js/type-aliases/BaseTokenizer.md | 19 +++ .../js/type-aliases/TokenizeTableOptions.md | 11 ++ nodejs/__test__/table.test.ts | 70 ++++++++ nodejs/lancedb/index.ts | 68 ++++++++ nodejs/lancedb/indices.ts | 16 +- nodejs/lancedb/table.ts | 44 +++++ nodejs/src/index.rs | 62 +++++++ nodejs/src/table.rs | 43 ++++- python/python/lancedb/__init__.py | 41 ++++- python/python/lancedb/_lancedb.pyi | 26 +++ python/python/lancedb/index.py | 2 + python/python/lancedb/remote/table.py | 18 ++ python/python/lancedb/table.py | 59 +++++++ python/python/lancedb/types.py | 4 +- python/python/tests/test_fts.py | 91 ++++++++++ python/src/lib.rs | 6 +- python/src/table.rs | 103 +++++++++++- rust/lancedb/src/lib.rs | 10 +- rust/lancedb/src/remote/table.rs | 138 ++++++++++++++- rust/lancedb/src/table.rs | 159 +++++++++++++++++- 26 files changed, 1175 insertions(+), 16 deletions(-) create mode 100644 docs/src/js/functions/tokenize.md create mode 100644 docs/src/js/interfaces/FtsToken.md create mode 100644 docs/src/js/interfaces/TokenizeOptions.md create mode 100644 docs/src/js/type-aliases/BaseTokenizer.md create mode 100644 docs/src/js/type-aliases/TokenizeTableOptions.md diff --git a/docs/src/js/classes/Table.md b/docs/src/js/classes/Table.md index 4e96fe9d3..05636657e 100644 --- a/docs/src/js/classes/Table.md +++ b/docs/src/js/classes/Table.md @@ -934,6 +934,32 @@ Return the table as an arrow table *** +### tokenize() + +```ts +abstract tokenize(query, options): Promise +``` + +Tokenize a full-text search query using the tokenizer configured on an FTS index. + +Specify exactly one of `column` or `indexName`. + +Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in +the client process from index metadata. For remote tables, this means the +same tokenizer model files must also exist locally. + +#### Parameters + +* **query**: `string` + +* **options**: [`TokenizeTableOptions`](../type-aliases/TokenizeTableOptions.md) + +#### Returns + +`Promise`<[`FtsToken`](../interfaces/FtsToken.md)[]> + +*** + ### unsetLsmWriteSpec() ```ts diff --git a/docs/src/js/functions/tokenize.md b/docs/src/js/functions/tokenize.md new file mode 100644 index 000000000..f31c0dda3 --- /dev/null +++ b/docs/src/js/functions/tokenize.md @@ -0,0 +1,26 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / tokenize + +# Function: tokenize() + +```ts +function tokenize(query, options?): Promise +``` + +Tokenize a full-text search query using an explicit tokenizer. + +This does not require a table or FTS index. The tokenizer options match +[Index.fts](../classes/Index.md#fts). + +## Parameters + +* **query**: `string` + +* **options?**: `Partial`<[`TokenizeOptions`](../interfaces/TokenizeOptions.md)> + +## Returns + +`Promise`<[`FtsToken`](../interfaces/FtsToken.md)[]> diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 7baaa4f92..907c89acc 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -72,6 +72,7 @@ - [FragmentStatistics](interfaces/FragmentStatistics.md) - [FragmentSummaryStats](interfaces/FragmentSummaryStats.md) - [FtsOptions](interfaces/FtsOptions.md) +- [FtsToken](interfaces/FtsToken.md) - [FullTextQuery](interfaces/FullTextQuery.md) - [FullTextSearchOptions](interfaces/FullTextSearchOptions.md) - [HnswPqOptions](interfaces/HnswPqOptions.md) @@ -107,6 +108,7 @@ - [TimeoutConfig](interfaces/TimeoutConfig.md) - [TlsConfig](interfaces/TlsConfig.md) - [TokenResponse](interfaces/TokenResponse.md) +- [TokenizeOptions](interfaces/TokenizeOptions.md) - [UpdateFieldMetadataResult](interfaces/UpdateFieldMetadataResult.md) - [UpdateOptions](interfaces/UpdateOptions.md) - [UpdateResult](interfaces/UpdateResult.md) @@ -116,6 +118,7 @@ ## Type Aliases +- [BaseTokenizer](type-aliases/BaseTokenizer.md) - [Data](type-aliases/Data.md) - [DataLike](type-aliases/DataLike.md) - [FieldLike](type-aliases/FieldLike.md) @@ -125,6 +128,7 @@ - [RecordBatchLike](type-aliases/RecordBatchLike.md) - [SchemaLike](type-aliases/SchemaLike.md) - [TableLike](type-aliases/TableLike.md) +- [TokenizeTableOptions](type-aliases/TokenizeTableOptions.md) ## Functions @@ -135,3 +139,4 @@ - [makeArrowTable](functions/makeArrowTable.md) - [packBits](functions/packBits.md) - [permutationBuilder](functions/permutationBuilder.md) +- [tokenize](functions/tokenize.md) diff --git a/docs/src/js/interfaces/FtsOptions.md b/docs/src/js/interfaces/FtsOptions.md index 0e9822163..12ab8fc84 100644 --- a/docs/src/js/interfaces/FtsOptions.md +++ b/docs/src/js/interfaces/FtsOptions.md @@ -23,7 +23,7 @@ whether to remove punctuation ### baseTokenizer? ```ts -optional baseTokenizer: "raw" | "simple" | "whitespace" | "ngram"; +optional baseTokenizer: BaseTokenizer; ``` The tokenizer to use when building the index. @@ -37,6 +37,10 @@ The following tokenizers are available: "raw" - Raw tokenizer. This tokenizer does not split the text into tokens and indexes the entire text as a single token. +"icu" - ICU dictionary-based word segmentation. + +"icu/split" - ICU segmentation with simple-style delimiter splitting. + *** ### language? diff --git a/docs/src/js/interfaces/FtsToken.md b/docs/src/js/interfaces/FtsToken.md new file mode 100644 index 000000000..338a8607a --- /dev/null +++ b/docs/src/js/interfaces/FtsToken.md @@ -0,0 +1,29 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / FtsToken + +# Interface: FtsToken + +Token produced by the tokenizer configured on a full-text search index. + +## Properties + +### position + +```ts +position: number; +``` + +Token position used by full-text query matching. + +*** + +### text + +```ts +text: string; +``` + +Token text after tokenizer filters have been applied. diff --git a/docs/src/js/interfaces/TokenizeOptions.md b/docs/src/js/interfaces/TokenizeOptions.md new file mode 100644 index 000000000..f061356d1 --- /dev/null +++ b/docs/src/js/interfaces/TokenizeOptions.md @@ -0,0 +1,109 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / TokenizeOptions + +# Interface: TokenizeOptions + +Options for tokenizing a full-text search query without a table index. + +## Properties + +### asciiFolding? + +```ts +optional asciiFolding: boolean; +``` + +Whether to fold ASCII characters. + +*** + +### baseTokenizer? + +```ts +optional baseTokenizer: BaseTokenizer; +``` + +The tokenizer to use. The default is "simple". + +*** + +### language? + +```ts +optional language: string; +``` + +Language for stemming and stop words. + +*** + +### lowercase? + +```ts +optional lowercase: boolean; +``` + +Whether to lowercase tokens. + +*** + +### maxTokenLength? + +```ts +optional maxTokenLength: number; +``` + +Maximum token length; tokens longer than this are ignored. + +*** + +### ngramMaxLength? + +```ts +optional ngramMaxLength: number; +``` + +N-gram maximum length. + +*** + +### ngramMinLength? + +```ts +optional ngramMinLength: number; +``` + +N-gram minimum length. + +*** + +### prefixOnly? + +```ts +optional prefixOnly: boolean; +``` + +Whether to only emit token prefixes for the n-gram tokenizer. + +*** + +### removeStopWords? + +```ts +optional removeStopWords: boolean; +``` + +Whether to remove stop words. + +*** + +### stem? + +```ts +optional stem: boolean; +``` + +Whether to stem tokens. diff --git a/docs/src/js/type-aliases/BaseTokenizer.md b/docs/src/js/type-aliases/BaseTokenizer.md new file mode 100644 index 000000000..75b377f27 --- /dev/null +++ b/docs/src/js/type-aliases/BaseTokenizer.md @@ -0,0 +1,19 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / BaseTokenizer + +# Type Alias: BaseTokenizer + +```ts +type BaseTokenizer: + | "simple" + | "whitespace" + | "raw" + | "ngram" + | "icu" + | "icu/split" + | `jieba/${string}` + | `lindera/${string}`; +``` diff --git a/docs/src/js/type-aliases/TokenizeTableOptions.md b/docs/src/js/type-aliases/TokenizeTableOptions.md new file mode 100644 index 000000000..5a641dc52 --- /dev/null +++ b/docs/src/js/type-aliases/TokenizeTableOptions.md @@ -0,0 +1,11 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / TokenizeTableOptions + +# Type Alias: TokenizeTableOptions + +```ts +type TokenizeTableOptions: object | object; +``` diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 925b0b155..dfe0af8c8 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -16,6 +16,7 @@ import { PhraseQuery, Table, connect, + tokenize, } from "../lancedb"; import { Table as ArrowTable, @@ -2307,6 +2308,75 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( expect(results2[0].text).toBe(data[1].text); }); + test("tokenizes FTS queries by column or index name", async () => { + const db = await connect(tmpDir.name); + const data = [ + { + text: "Running in cafés", + japanese: "Hello, こんにちは世界!", + vector: [0.1, 0.2, 0.3], + }, + ]; + const table = await db.createTable("test", data); + await table.createIndex("text", { + config: Index.fts({ baseTokenizer: "simple" }), + }); + await table.createIndex("japanese", { + config: Index.fts({ + baseTokenizer: "icu", + stem: false, + removeStopWords: false, + }), + name: "japanese_icu_idx", + }); + + await expect(table.tokenize("hello", {} as never)).rejects.toThrow( + "Specify exactly one", + ); + await expect( + table.tokenize("hello", { + column: "text", + indexName: "text_idx", + } as never), + ).rejects.toThrow("Specify exactly one"); + + const simpleTokens = await table.tokenize("Running in cafés", { + column: "text", + }); + expect(simpleTokens).toEqual([ + { text: "run", position: 0 }, + { text: "cafe", position: 2 }, + ]); + + const icuTokens = await table.tokenize("Hello, こんにちは世界!", { + indexName: "japanese_icu_idx", + }); + expect(icuTokens).toEqual([ + { text: "hello", position: 0 }, + { text: "こんにちは", position: 1 }, + { text: "世界", position: 2 }, + ]); + + const directSimpleTokens = await tokenize("Running in cafés", { + baseTokenizer: "simple", + }); + expect(directSimpleTokens).toEqual([ + { text: "run", position: 0 }, + { text: "cafe", position: 2 }, + ]); + + const directIcuTokens = await tokenize("Hello, こんにちは世界!", { + baseTokenizer: "icu", + stem: false, + removeStopWords: false, + }); + expect(directIcuTokens).toEqual([ + { text: "hello", position: 0 }, + { text: "こんにちは", position: 1 }, + { text: "世界", position: 2 }, + ]); + }); + test("full text search fast search", async () => { const db = await connect(tmpDir.name); const data = [{ text: "hello world", vector: [0.1, 0.2, 0.3], id: 1 }]; diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index ebd638634..527f5f4a2 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -13,9 +13,12 @@ import { Connection as LanceDbConnection, JsHeaderProvider as NativeJsHeaderProvider, Session, + tokenize as nativeTokenize, } from "./native.js"; import { HeaderProvider } from "./header"; +import type { BaseTokenizer } from "./indices"; +import type { FtsToken } from "./table"; // Re-export native header provider for use with connectWithHeaderProvider export { JsHeaderProvider as NativeJsHeaderProvider } from "./native.js"; @@ -114,6 +117,7 @@ export { HnswPqOptions, HnswSqOptions, FtsOptions, + BaseTokenizer, } from "./indices"; export { @@ -124,6 +128,8 @@ export { OptimizeOptions, Version, WriteProgress, + FtsToken, + TokenizeTableOptions, LsmWriteSpec, ColumnAlteration, FieldMetadataUpdate, @@ -155,6 +161,68 @@ export { } from "./arrow"; export { IntoSql, packBits } from "./util"; +/** + * Options for tokenizing a full-text search query without a table index. + */ +export interface TokenizeOptions { + /** + * The tokenizer to use. The default is "simple". + */ + baseTokenizer?: BaseTokenizer; + + /** Language for stemming and stop words. */ + language?: string; + + /** Maximum token length; tokens longer than this are ignored. */ + maxTokenLength?: number; + + /** Whether to lowercase tokens. */ + lowercase?: boolean; + + /** Whether to stem tokens. */ + stem?: boolean; + + /** Whether to remove stop words. */ + removeStopWords?: boolean; + + /** Whether to fold ASCII characters. */ + asciiFolding?: boolean; + + /** N-gram minimum length. */ + ngramMinLength?: number; + + /** N-gram maximum length. */ + ngramMaxLength?: number; + + /** Whether to only emit token prefixes for the n-gram tokenizer. */ + prefixOnly?: boolean; +} + +/** + * Tokenize a full-text search query using an explicit tokenizer. + * + * This does not require a table or FTS index. The tokenizer options match + * {@link Index.fts}. + */ +export async function tokenize( + query: string, + options?: Partial, +): Promise { + return await nativeTokenize( + query, + options?.baseTokenizer, + options?.language, + options?.maxTokenLength, + options?.lowercase, + options?.stem, + options?.removeStopWords, + options?.asciiFolding, + options?.ngramMinLength, + options?.ngramMaxLength, + options?.prefixOnly, + ); +} + /** * Connect to a LanceDB instance at the given URI. * diff --git a/nodejs/lancedb/indices.ts b/nodejs/lancedb/indices.ts index c82c3edb8..e45556e5f 100644 --- a/nodejs/lancedb/indices.ts +++ b/nodejs/lancedb/indices.ts @@ -486,6 +486,16 @@ export interface IvfFlatOptions { sampleRate?: number; } +export type BaseTokenizer = + | "simple" + | "whitespace" + | "raw" + | "ngram" + | "icu" + | "icu/split" + | `jieba/${string}` + | `lindera/${string}`; + /** * Options to create a full text search index */ @@ -509,8 +519,12 @@ export interface FtsOptions { * "whitespace" - Whitespace tokenizer. This tokenizer splits the text into tokens using whitespace as a delimiter. * * "raw" - Raw tokenizer. This tokenizer does not split the text into tokens and indexes the entire text as a single token. + * + * "icu" - ICU dictionary-based word segmentation. + * + * "icu/split" - ICU segmentation with simple-style delimiter splitting. */ - baseTokenizer?: "simple" | "whitespace" | "raw" | "ngram"; + baseTokenizer?: BaseTokenizer; /** * language for stemming and stop words diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index 84318e63c..07222a192 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -158,6 +158,26 @@ export interface Version { metadata: Record; } +/** Token produced by the tokenizer configured on a full-text search index. */ +export interface FtsToken { + /** Token text after tokenizer filters have been applied. */ + text: string; + /** Token position used by full-text query matching. */ + position: number; +} + +export type TokenizeTableOptions = + | { + /** FTS-indexed column whose tokenizer should be used. */ + column: string; + indexName?: never; + } + | { + /** Name of the FTS index whose tokenizer should be used. */ + indexName: string; + column?: never; + }; + /** * Specification selecting Lance's MemWAL LSM-style write path for * `mergeInsert`. @@ -716,6 +736,19 @@ export abstract class Table { abstract optimize(options?: Partial): Promise; /** List all indices that have been created with {@link Table.createIndex} */ abstract listIndices(): Promise; + /** + * Tokenize a full-text search query using the tokenizer configured on an FTS index. + * + * Specify exactly one of `column` or `indexName`. + * + * Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in + * the client process from index metadata. For remote tables, this means the + * same tokenizer model files must also exist locally. + */ + abstract tokenize( + query: string, + options: TokenizeTableOptions, + ): Promise; /** Return the table as an arrow table */ abstract toArrow(): Promise; @@ -1173,6 +1206,17 @@ export class LocalTable extends Table { return await this.inner.listIndices(); } + async tokenize( + query: string, + options: TokenizeTableOptions, + ): Promise { + return await this.inner.tokenize( + query, + options?.column, + options?.indexName, + ); + } + async toArrow(): Promise { return await this.query().toArrow(); } diff --git a/nodejs/src/index.rs b/nodejs/src/index.rs index 868f15e89..db84710ba 100644 --- a/nodejs/src/index.rs +++ b/nodejs/src/index.rs @@ -9,8 +9,11 @@ use lancedb::index::vector::{ IvfFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder, IvfRqIndexBuilder, }; +use lancedb::tokenize as lancedb_tokenize; use napi_derive::napi; +use crate::error::NapiErrorExt; +use crate::table::FtsToken; use crate::util::parse_distance_type; #[napi] @@ -30,6 +33,65 @@ impl Index { } } +#[napi(catch_unwind)] +#[allow(dead_code, clippy::too_many_arguments)] +pub fn tokenize( + query: String, + base_tokenizer: Option, + language: Option, + max_token_length: Option, + lower_case: Option, + stem: Option, + remove_stop_words: Option, + ascii_folding: Option, + ngram_min_length: Option, + ngram_max_length: Option, + prefix_only: Option, +) -> napi::Result> { + let mut opts = FtsIndexBuilder::default(); + if let Some(base_tokenizer) = base_tokenizer { + opts = opts.base_tokenizer(base_tokenizer); + } + if let Some(language) = language { + opts = opts.language(&language).map_err(|_| { + napi::Error::from_reason(format!( + "LanceDB does not support the requested language: '{}'", + language + )) + })?; + } + if let Some(max_token_length) = max_token_length { + opts = opts.max_token_length(Some(max_token_length as usize)); + } + if let Some(lower_case) = lower_case { + opts = opts.lower_case(lower_case); + } + if let Some(stem) = stem { + opts = opts.stem(stem); + } + if let Some(remove_stop_words) = remove_stop_words { + opts = opts.remove_stop_words(remove_stop_words); + } + if let Some(ascii_folding) = ascii_folding { + opts = opts.ascii_folding(ascii_folding); + } + if let Some(ngram_min_length) = ngram_min_length { + opts = opts.ngram_min_length(ngram_min_length); + } + if let Some(ngram_max_length) = ngram_max_length { + opts = opts.ngram_max_length(ngram_max_length); + } + if let Some(prefix_only) = prefix_only { + opts = opts.ngram_prefix_only(prefix_only); + } + + Ok(lancedb_tokenize(&query, &opts) + .default_error()? + .into_iter() + .map(FtsToken::from) + .collect()) +} + #[napi] impl Index { #[napi(factory)] diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index 251282d88..8252a8c3b 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -8,8 +8,8 @@ use chrono::{DateTime, Utc}; use lancedb::ipc::{ipc_file_to_batches, ipc_file_to_schema}; use lancedb::table::{ AddDataMode, ColumnAlteration as LanceColumnAlteration, Duration, - FieldMetadataUpdate as LanceFieldMetadataUpdate, NewColumnTransform, OptimizeAction, - OptimizeOptions, Ref, Table as LanceDbTable, + FieldMetadataUpdate as LanceFieldMetadataUpdate, FtsToken as LanceDbFtsToken, + NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable, }; use napi::bindgen_prelude::*; use napi::threadsafe_function::{ThreadsafeFunction, ThreadsafeFunctionCallMode}; @@ -574,6 +574,27 @@ impl Table { .collect::>()) } + #[napi(catch_unwind)] + pub async fn tokenize( + &self, + query: String, + column: Option, + index_name: Option, + ) -> napi::Result> { + let table = self.inner_ref()?; + let tokens = match (column.as_deref(), index_name.as_deref()) { + (Some(_), Some(_)) | (None, None) => { + return Err(napi::Error::from_reason( + "Specify exactly one of 'column' or 'indexName'", + )); + } + (Some(column), None) => table.tokenize_with_column(&query, column).await, + (None, Some(index_name)) => table.tokenize(&query, index_name).await, + } + .default_error()?; + Ok(tokens.into_iter().map(FtsToken::from).collect()) + } + #[napi(catch_unwind)] pub async fn index_stats(&self, index_name: String) -> napi::Result> { let tbl = self.inner_ref()?; @@ -681,6 +702,24 @@ impl From for IndexConfig { } } +#[napi(object)] +/// A token produced by the tokenizer configured on a full-text search index. +pub struct FtsToken { + /// The token text after the index tokenizer has applied its filters. + pub text: String, + /// The token position used by full-text query matching. + pub position: u32, +} + +impl From for FtsToken { + fn from(token: LanceDbFtsToken) -> Self { + Self { + text: token.text, + position: token.position, + } + } +} + /// Specification selecting Lance's MemWAL LSM-style write path for /// `mergeInsert`. /// diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 6cdd6e855..6f6468d9c 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -6,11 +6,13 @@ import importlib.metadata import os from concurrent.futures import ThreadPoolExecutor from datetime import timedelta -from typing import Dict, Optional, Union, Any, List +from typing import Dict, Optional, Union, Any, List, Iterable __version__ = importlib.metadata.version("lancedb") from ._lancedb import connect as lancedb_connect +from ._lancedb import FtsToken +from ._lancedb import tokenize as _tokenize from .common import URI, sanitize_uri from urllib.parse import urlparse from .db import AsyncConnection, DBConnection, LanceDBConnection @@ -19,6 +21,7 @@ from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func from .schema import blob, vector, BlobType from .table import AsyncTable, Table +from .types import BaseTokenizerType from ._lancedb import Session from .namespace import ( connect_namespace, @@ -246,6 +249,40 @@ def connect( ) +def tokenize( + query: str, + *, + base_tokenizer: BaseTokenizerType = "simple", + language: str = "English", + max_token_length: Optional[int] = 40, + lower_case: bool = True, + stem: bool = True, + remove_stop_words: bool = True, + ascii_folding: bool = True, + ngram_min_length: int = 3, + ngram_max_length: int = 3, + prefix_only: bool = False, +) -> Iterable[FtsToken]: + """Tokenize a full-text search query using an explicit tokenizer. + + This does not require a table or FTS index. The tokenizer options match + :class:`lancedb.index.FTS`. + """ + return _tokenize( + query, + base_tokenizer=base_tokenizer, + language=language, + max_token_length=max_token_length, + lower_case=lower_case, + stem=stem, + remove_stop_words=remove_stop_words, + ascii_folding=ascii_folding, + ngram_min_length=ngram_min_length, + ngram_max_length=ngram_max_length, + prefix_only=prefix_only, + ) + + WORKER_PROPERTY_PREFIX = "_lancedb_worker_" @@ -456,11 +493,13 @@ async def connect_async( __all__ = [ "connect", "connect_async", + "tokenize", "connect_namespace", "connect_namespace_async", "AsyncConnection", "AsyncLanceNamespaceDBConnection", "AsyncTable", + "FtsToken", "col", "Expr", "func", diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 390235d8c..368bdd069 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -25,6 +25,7 @@ from lance_namespace import ( ListTablesResponse, ) from .remote import ClientConfig +from .types import BaseTokenizerType IvfHnswPq: type[HnswPq] = HnswPq IvfHnswSq: type[HnswSq] = HnswSq @@ -48,6 +49,20 @@ class MetricDescription: def register_lancedb_metrics_recorder() -> bool: ... def lancedb_metrics_catalog() -> List[MetricDescription]: ... def snapshot_lancedb_metrics() -> List[MetricPoint]: ... +def tokenize( + query: str, + *, + base_tokenizer: BaseTokenizerType = "simple", + language: str = "English", + max_token_length: Optional[int] = 40, + lower_case: bool = True, + stem: bool = True, + remove_stop_words: bool = True, + ascii_folding: bool = True, + ngram_min_length: int = 3, + ngram_max_length: int = 3, + prefix_only: bool = False, +) -> List["FtsToken"]: ... class PyExpr: """A type-safe DataFusion expression node (Rust-side handle).""" @@ -238,6 +253,13 @@ class Table: async def prewarm_index(self, index_name: str) -> None: ... async def prewarm_data(self, columns: Optional[List[str]] = None) -> None: ... async def list_indices(self) -> list[IndexConfig]: ... + async def tokenize( + self, + query: str, + *, + column: Optional[str] = None, + index_name: Optional[str] = None, + ) -> list[FtsToken]: ... async def delete(self, filter: Union[str, PyExpr]) -> DeleteResult: ... async def add_columns(self, columns: list[tuple[str, str]]) -> AddColumnsResult: ... async def add_columns_with_schema(self, schema: pa.Schema) -> AddColumnsResult: ... @@ -511,6 +533,10 @@ class MergeResult: num_attempts: int num_rows: int +class FtsToken: + text: str + position: int + class LsmWriteSpec: """Specification selecting Lance's MemWAL LSM-style write path for `merge_insert`.""" diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index 0fb6d45ba..76daaed11 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -127,6 +127,8 @@ class FTS: - "whitespace": Split text by whitespace, but not punctuation. - "raw": No tokenization. The entire text is treated as a single token. - "ngram": N-gram tokenizer for substring-style matching. + - "icu": ICU dictionary-based word segmentation. + - "icu/split": ICU segmentation with simple-style delimiter splitting. - "jieba/*": Jieba tokenizer loaded from Lance's language model home. - "lindera/*": Lindera tokenizer loaded from Lance's language model home. language : str, default "English" diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 772e0055e..75ad59b50 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -28,6 +28,7 @@ from lancedb._lancedb import ( UpdateFieldMetadataResult, DeleteResult, DropColumnsResult, + FtsToken, IndexConfig, LsmWriteSpec, MergeResult, @@ -244,6 +245,23 @@ class RemoteTable(Table): """List all the indices on the table""" return LOOP.run(self._table.list_indices()) + def tokenize( + self, + query: str, + *, + column: Optional[str] = None, + index_name: Optional[str] = None, + ) -> Iterable[FtsToken]: + """Tokenize a query using the tokenizer configured on an FTS index. + + Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are + rebuilt in the client process from index metadata, so the same tokenizer + model files must exist locally. + """ + return LOOP.run( + self._table.tokenize(query, column=column, index_name=index_name) + ) + def index_stats(self, index_uuid: str) -> Optional[IndexStatistics]: """List all the stats of a specified index""" return LOOP.run(self._table.index_stats(index_uuid)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 98f14dbfe..8243c06cf 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -173,6 +173,7 @@ if TYPE_CHECKING: UpdateFieldMetadataResult, DeleteResult, DropColumnsResult, + FtsToken, LsmWriteSpec, MergeResult, UpdateResult, @@ -1147,6 +1148,8 @@ class Table(ABC): - "whitespace": Split text by whitespace, but not punctuation. - "raw": No tokenization. The entire text is treated as a single token. - "ngram": N-Gram tokenizer. + - "icu": ICU dictionary-based word segmentation. + - "icu/split": ICU segmentation with simple-style delimiter splitting. - "jieba/*": Jieba tokenizer loaded from Lance's language model home. - "lindera/*": Lindera tokenizer loaded from Lance's language model home. language : str, default "English" @@ -1799,6 +1802,24 @@ class Table(ABC): [Table.create_index][lancedb.table.Table.create_index] """ + @abstractmethod + def tokenize( + self, + query: str, + *, + column: Optional[str] = None, + index_name: Optional[str] = None, + ) -> Iterable[FtsToken]: + """ + Tokenize a query using the tokenizer configured on an FTS index. + + Specify exactly one of ``column`` or ``index_name``. + + Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are + rebuilt in the client process from index metadata. For remote tables, + this means the same tokenizer model files must also exist locally. + """ + @abstractmethod def index_stats(self, index_name: str) -> Optional[IndexStatistics]: """ @@ -3744,6 +3765,26 @@ class LanceTable(Table): """ return LOOP.run(self._table.list_indices()) + def tokenize( + self, + query: str, + *, + column: Optional[str] = None, + index_name: Optional[str] = None, + ) -> Iterable[FtsToken]: + """ + Tokenize a query using the tokenizer configured on an FTS index. + + Specify exactly one of ``column`` or ``index_name``. + + Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are + rebuilt in the client process from index metadata. For remote tables, + this means the same tokenizer model files must also exist locally. + """ + return LOOP.run( + self._table.tokenize(query, column=column, index_name=index_name) + ) + def index_stats(self, index_name: str) -> Optional[IndexStatistics]: """ Retrieve statistics about an index @@ -5804,6 +5845,24 @@ class AsyncTable: """ return await self._inner.list_indices() + async def tokenize( + self, + query: str, + *, + column: Optional[str] = None, + index_name: Optional[str] = None, + ) -> Iterable[FtsToken]: + """ + Tokenize a query using the tokenizer configured on an FTS index. + + Specify exactly one of ``column`` or ``index_name``. + + Model-backed tokenizers such as ``jieba/*`` and ``lindera/*`` are + rebuilt in the client process from index metadata. For remote tables, + this means the same tokenizer model files must also exist locally. + """ + return await self._inner.tokenize(query, column=column, index_name=index_name) + async def index_stats(self, index_name: str) -> Optional[IndexStatistics]: """ Retrieve statistics about an index diff --git a/python/python/lancedb/types.py b/python/python/lancedb/types.py index d2edd318a..567eecb4e 100644 --- a/python/python/lancedb/types.py +++ b/python/python/lancedb/types.py @@ -55,5 +55,7 @@ IndexType = Literal[ ] # Tokenizer literals -BuiltinTokenizerType = Literal["simple", "raw", "whitespace", "ngram"] +BuiltinTokenizerType = Literal[ + "simple", "raw", "whitespace", "ngram", "icu", "icu/split" +] BaseTokenizerType = BuiltinTokenizerType | str diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index 134b27980..f822d1692 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -786,6 +786,97 @@ def test_language(mem_db: DBConnection): assert len(results) == 0 +def test_tokenize_uses_simple_index_tokenizer(mem_db: DBConnection): + data = pa.table({"text": ["Running in cafés"], "other": ["Running in cafés"]}) + table = mem_db.create_table("test_tokenize", data=data) + table.create_index("text", config=FTS(base_tokenizer="simple")) + + tokens = table.tokenize("Running in cafés", column="text") + + assert [(token.text, token.position) for token in tokens] == [ + ("run", 0), + ("cafe", 2), + ] + + +def test_tokenize_uses_icu_index_tokenizer_by_name(mem_db: DBConnection): + data = pa.table({"text": ["Hello, こんにちは世界!"]}) + table = mem_db.create_table("test_tokenize_icu", data=data) + table.create_index( + "text", + config=FTS( + base_tokenizer="icu", + stem=False, + remove_stop_words=False, + ), + name="text_icu_idx", + ) + + tokens = table.tokenize("Hello, こんにちは世界!", index_name="text_icu_idx") + + assert [(token.text, token.position) for token in tokens] == [ + ("hello", 0), + ("こんにちは", 1), + ("世界", 2), + ] + + +def test_tokenize_requires_one_selector(mem_db: DBConnection): + data = pa.table({"text": ["hello world"]}) + table = mem_db.create_table("test_tokenize_selector", data=data) + table.create_index("text", config=FTS(), name="text_idx") + + with pytest.raises(ValueError, match="Specify exactly one"): + table.tokenize("hello") + + with pytest.raises(ValueError, match="Specify exactly one"): + table.tokenize("hello", column="text", index_name="text_idx") + + +def test_tokenize_requires_fts_index(mem_db: DBConnection): + data = pa.table({"text": ["hello world"]}) + table = mem_db.create_table("test_tokenize_no_index", data=data) + + with pytest.raises(ValueError, match="does not have a full text search index"): + table.tokenize("hello", column="text") + + +@pytest.mark.asyncio +async def test_tokenize_async(async_table): + await async_table.create_index("text", config=FTS()) + + tokens = await async_table.tokenize("Running in cafés", column="text") + + assert [(token.text, token.position) for token in tokens] == [ + ("run", 0), + ("cafe", 2), + ] + + +def test_tokenize_uses_explicit_simple_tokenizer(): + tokens = ldb.tokenize("Running in cafés", base_tokenizer="simple") + + assert [(token.text, token.position) for token in tokens] == [ + ("run", 0), + ("cafe", 2), + ] + + +def test_tokenize_uses_explicit_icu_tokenizer(): + tokens = ldb.tokenize( + "Hello, こんにちは世界!", + base_tokenizer="icu", + stem=False, + remove_stop_words=False, + ) + + assert [(token.text, token.position) for token in tokens] == [ + ("hello", 0), + ("こんにちは", 1), + ("世界", 2), + ] + + def test_fts_on_list(mem_db: DBConnection): data = pa.table( { diff --git a/python/src/lib.rs b/python/src/lib.rs index a72386305..d480f9613 100644 --- a/python/src/lib.rs +++ b/python/src/lib.rs @@ -15,8 +15,8 @@ use pyo3::{ use query::{FTSQuery, HybridQuery, Query, VectorQuery}; use session::Session; use table::{ - AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, LsmWriteSpec, - MergeResult, PyBlobFile, Table, UpdateFieldMetadataResult, UpdateResult, + AddColumnsResult, AddResult, AlterColumnsResult, DeleteResult, DropColumnsResult, FtsToken, + LsmWriteSpec, MergeResult, PyBlobFile, Table, UpdateFieldMetadataResult, UpdateResult, }; pub mod arrow; @@ -60,6 +60,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; @@ -75,6 +76,7 @@ pub fn _lancedb(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(connect, m)?)?; m.add_function(wrap_pyfunction!(connect_namespace, m)?)?; m.add_function(wrap_pyfunction!(connect_namespace_client, m)?)?; + m.add_function(wrap_pyfunction!(table::tokenize, m)?)?; m.add_function(wrap_pyfunction!(permutation::async_permutation_builder, m)?)?; m.add_function(wrap_pyfunction!(util::validate_table_name, m)?)?; m.add_function(wrap_pyfunction!(query::fts_query_to_json, m)?)?; diff --git a/python/src/table.rs b/python/src/table.rs index fe2504ec2..c9784facb 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -18,14 +18,16 @@ use arrow::{ pyarrow::{FromPyArrow, PyArrowType, ToPyArrow}, }; use lancedb::blob::BlobFile; +use lancedb::index::scalar::FtsIndexBuilder; use lancedb::table::{ - AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, NewColumnTransform, - OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable, + AddDataMode, ColumnAlteration, Duration, FieldMetadataUpdate, FtsToken as LanceDbFtsToken, + NewColumnTransform, OptimizeAction, OptimizeOptions, Ref, Table as LanceDbTable, }; +use lancedb::tokenize as lancedb_tokenize; use pyo3::{ Bound, FromPyObject, Py, PyAny, PyRef, PyResult, Python, exceptions::{PyRuntimeError, PyValueError}, - pyclass, pymethods, + pyclass, pyfunction, pymethods, types::{IntoPyDict, PyAnyMethods, PyBytes, PyDict, PyDictMethods}, }; @@ -486,6 +488,78 @@ impl PyBlobFile { } } +#[pyclass(get_all, from_py_object)] +#[derive(Clone, Debug)] +pub struct FtsToken { + pub text: String, + pub position: u32, +} + +#[pymethods] +impl FtsToken { + pub fn __repr__(&self) -> String { + format!("FtsToken(text={:?}, position={})", self.text, self.position) + } +} + +impl From for FtsToken { + fn from(token: LanceDbFtsToken) -> Self { + Self { + text: token.text, + position: token.position, + } + } +} + +#[pyfunction(signature = ( + query, + *, + base_tokenizer = "simple".to_string(), + language = "English".to_string(), + max_token_length = Some(40), + lower_case = true, + stem = true, + remove_stop_words = true, + ascii_folding = true, + ngram_min_length = 3, + ngram_max_length = 3, + prefix_only = false +))] +#[allow(clippy::too_many_arguments)] +pub fn tokenize( + query: String, + base_tokenizer: String, + language: String, + max_token_length: Option, + lower_case: bool, + stem: bool, + remove_stop_words: bool, + ascii_folding: bool, + ngram_min_length: u32, + ngram_max_length: u32, + prefix_only: bool, +) -> PyResult> { + let params = FtsIndexBuilder::default() + .base_tokenizer(base_tokenizer) + .language(&language) + .map_err(|_| { + PyValueError::new_err(format!( + "LanceDB does not support the requested language: '{}'", + language + )) + })? + .max_token_length(max_token_length.map(|value| value as usize)) + .lower_case(lower_case) + .stem(stem) + .remove_stop_words(remove_stop_words) + .ascii_folding(ascii_folding) + .ngram_min_length(ngram_min_length) + .ngram_max_length(ngram_max_length) + .ngram_prefix_only(prefix_only); + let tokens = lancedb_tokenize(&query, ¶ms).infer_error()?; + Ok(tokens.into_iter().map(FtsToken::from).collect()) +} + #[pyclass] pub struct Table { // We keep a copy of the name to use if the inner table is dropped @@ -784,6 +858,29 @@ impl Table { }) } + #[pyo3(signature = (query, *, column=None, index_name=None))] + pub fn tokenize( + self_: PyRef<'_, Self>, + query: String, + column: Option, + index_name: Option, + ) -> PyResult> { + let inner = self_.inner_ref()?.clone(); + future_into_py(self_.py(), async move { + let tokens = match (column.as_deref(), index_name.as_deref()) { + (Some(_), Some(_)) | (None, None) => { + return Err(PyValueError::new_err( + "Specify exactly one of 'column' or 'index_name'", + )); + } + (Some(column), None) => inner.tokenize_with_column(&query, column).await, + (None, Some(index_name)) => inner.tokenize(&query, index_name).await, + } + .infer_error()?; + Ok(tokens.into_iter().map(FtsToken::from).collect::>()) + }) + } + pub fn index_stats(self_: PyRef<'_, Self>, index_name: String) -> PyResult> { let inner = self_.inner_ref()?.clone(); future_into_py(self_.py(), async move { diff --git a/rust/lancedb/src/lib.rs b/rust/lancedb/src/lib.rs index f4100ae84..0b2962682 100644 --- a/rust/lancedb/src/lib.rs +++ b/rust/lancedb/src/lib.rs @@ -207,7 +207,15 @@ use lance_linalg::distance::DistanceType as LanceDistanceType; /// a built-in pull-based adapter. #[cfg(feature = "metrics")] pub use metrics; -pub use table::Table; +pub use table::{FtsToken, Table}; + +/// Tokenize a full-text search query using an explicit FTS tokenizer configuration. +/// +/// This does not require a table or FTS index. The tokenizer options are the +/// same [`index::scalar::FtsIndexBuilder`] values used when creating an FTS index. +pub fn tokenize(query: &str, params: &index::scalar::FtsIndexBuilder) -> Result> { + table::tokenize(query, params) +} #[derive(Debug, Copy, Clone, PartialEq, Serialize, Deserialize, Default)] #[non_exhaustive] diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index c9ec28ee4..3aa4246e9 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2817,8 +2817,7 @@ mod tests { use super::*; use crate::remote::client::{ClientConfig, RetryConfig}; - use crate::table::AddDataMode; - use crate::table::FieldMetadataUpdate; + use crate::table::{AddDataMode, FieldMetadataUpdate, FtsToken}; use arrow::{array::AsArray, compute::concat_batches, datatypes::Int32Type}; use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, record_batch}; @@ -4888,6 +4887,141 @@ mod tests { assert_eq!(text_idx.created_at, None); } + #[tokio::test] + async fn test_tokenize_uses_remote_index_details() { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let index_details = serde_json::json!({ + "base_tokenizer": "icu", + "language": "English", + "with_position": false, + "max_token_length": 40, + "lower_case": true, + "stem": false, + "remove_stop_words": false, + "ascii_folding": true, + }) + .to_string(); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + "index_details": index_details, + }, + ] + }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let tokens = table + .tokenize("Hello, こんにちは世界!", "text_idx") + .await + .unwrap(); + + assert_eq!( + tokens, + vec![ + FtsToken { + text: "hello".to_string(), + position: 0, + }, + FtsToken { + text: "こんにちは".to_string(), + position: 1, + }, + FtsToken { + text: "世界".to_string(), + position: 2, + }, + ] + ); + } + + #[tokio::test] + async fn test_tokenize_requires_existing_index_name() { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let table = Table::new_with_handler("my_table", move |request| -> http::Response { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ "indexes": [] }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let err = table.tokenize("hello", "text_idx").await.unwrap_err(); + assert!(matches!( + err, + Error::InvalidInput { message } + if message.contains("No index named 'text_idx'") + )); + } + + #[tokio::test] + async fn test_tokenize_with_column_remote_requires_index_details() { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let table = Table::new_with_handler("my_table", move |request| { + assert_eq!(request.method(), "POST"); + match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap(), + "/v1/table/my_table/index/list/" => { + let body = serde_json::json!({ + "indexes": [ + { + "index_name": "text_idx", + "columns": ["text"], + "index_type": "FTS", + }, + ] + }); + http::Response::builder() + .status(200) + .body(serde_json::to_string(&body).unwrap()) + .unwrap() + } + path => panic!("Unexpected path: {}", path), + } + }); + + let err = table + .tokenize_with_column("hello", "text") + .await + .unwrap_err(); + + assert!(matches!( + err, + Error::InvalidInput { message } + if message.contains("does not include tokenizer details") + )); + } + #[test] fn test_deserialize_created_at() { #[derive(Deserialize)] diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index d64b9132d..580b97010 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -28,6 +28,8 @@ use lance_io::object_store::{LanceNamespaceStorageOptionsProvider, StorageOption pub use query::AnyQuery; use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore; +use lance_index::scalar::InvertedIndexParams; +use lance_index::scalar::inverted::query::collect_query_tokens; use lance_namespace::LanceNamespace; use lance_namespace::error::NamespaceError; use lance_namespace::models::DescribeTableRequest; @@ -51,10 +53,10 @@ use crate::embeddings::{EmbeddingDefinition, EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; use crate::index::IndexStatistics; use crate::index::{Index, IndexBuilder}; -use crate::index::{IndexConfig, IndexStatisticsImpl}; +use crate::index::{IndexConfig, IndexStatisticsImpl, IndexType}; use crate::query::{IntoQueryVector, Query, QueryExecutionOptions, TakeQuery, VectorQuery}; use crate::table::datafusion::insert::InsertExec; -use crate::utils::{PatchReadParam, PatchWriteParam}; +use crate::utils::{PatchReadParam, PatchWriteParam, resolve_arrow_field_path}; use self::dataset::DatasetConsistencyWrapper; use self::merge::MergeInsertBuilder; @@ -473,6 +475,33 @@ impl LsmWriteSpec { } } +/// A token produced by the tokenizer configured on a full-text search index. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct FtsToken { + /// The token text after the index tokenizer has applied its filters. + pub text: String, + /// The token position used by full-text query matching. + pub position: u32, +} + +/// Tokenize a full-text search query using an explicit FTS tokenizer configuration. +/// +/// This does not require a table or FTS index. Use +/// [`crate::index::scalar::FtsIndexBuilder`] to supply the same tokenizer +/// options used when creating an FTS index. +pub fn tokenize(query: &str, params: &InvertedIndexParams) -> Result> { + let mut tokenizer = params.build().map_err(|err| Error::InvalidInput { + message: format!("Failed to build tokenizer: {}", err), + })?; + let tokens = collect_query_tokens(query, &mut tokenizer); + Ok((0..tokens.len()) + .map(|idx| FtsToken { + text: tokens.get_token(idx).to_string(), + position: tokens.position(idx), + }) + .collect()) +} + /// A trait for anything "table-like". This is used for both native tables (which target /// Lance datasets) and remote tables (which target LanceDB cloud) /// @@ -1661,6 +1690,111 @@ impl Table { self.inner.list_indices().await } + /// Tokenize a full-text search query using the tokenizer configured on an FTS index. + /// + /// Model-backed tokenizers such as `jieba/*` and `lindera/*` are rebuilt in + /// the client process from index metadata. For remote tables, this means the + /// same tokenizer model files must also exist locally. + pub async fn tokenize(&self, query: &str, index_name: &str) -> Result> { + let indices = self.inner.list_indices().await?; + let matches = indices + .iter() + .filter(|idx| idx.name == index_name) + .collect::>(); + let index = match matches.as_slice() { + [index] => *index, + [] => { + return Err(Error::InvalidInput { + message: format!("No index named '{}'", index_name), + }); + } + _ => { + return Err(Error::InvalidInput { + message: format!("Index name '{}' is ambiguous", index_name), + }); + } + }; + if index.index_type != IndexType::FTS { + return Err(Error::InvalidInput { + message: format!("Index '{}' is not a full text search index", index_name), + }); + } + self.tokenize_with_index(query, index, index_name) + } + + /// Tokenize a full-text search query using the tokenizer configured on the + /// FTS index for a column. + /// + /// The column must have exactly one FTS index. Model-backed tokenizers such + /// as `jieba/*` and `lindera/*` are rebuilt in the client process from + /// index metadata. For remote tables, this means the same tokenizer model + /// files must also exist locally. + pub async fn tokenize_with_column(&self, query: &str, column: &str) -> Result> { + let schema = self.inner.schema().await?; + let (column, _) = resolve_arrow_field_path(schema.as_ref(), column)?; + let indices = self.inner.list_indices().await?; + let matches = indices + .iter() + .filter(|idx| { + idx.index_type == IndexType::FTS + && idx.columns.len() == 1 + && idx.columns[0] == column + }) + .collect::>(); + let index = match matches.as_slice() { + [index] => *index, + [] => { + return Err(Error::InvalidInput { + message: format!("Column '{}' does not have a full text search index", column), + }); + } + _ => { + return Err(Error::InvalidInput { + message: format!( + "Column '{}' has multiple full text search indexes; tokenization by column is ambiguous", + column + ), + }); + } + }; + self.tokenize(query, &index.name).await + } + + fn tokenize_with_index( + &self, + query: &str, + index: &IndexConfig, + index_name: &str, + ) -> Result> { + let selector_description = format!("index name '{}'", index_name); + let details = index + .index_details + .as_deref() + .ok_or_else(|| Error::InvalidInput { + message: format!( + "Full text search index '{}' for {} does not include tokenizer details", + index.name, selector_description + ), + })?; + let params = serde_json::from_str::(details).map_err(|err| { + Error::InvalidInput { + message: format!( + "Failed to parse tokenizer details for full text search index '{}' for {}: {}", + index.name, selector_description, err + ), + } + })?; + tokenize(query, ¶ms).map_err(|err| match err { + Error::InvalidInput { message } => Error::InvalidInput { + message: format!( + "{} for full text search index '{}' for {}", + message, index.name, selector_description + ), + }, + err => err, + }) + } + /// Get the table URI (storage location) /// /// Returns the full storage location of the table (e.g., S3/GCS path). @@ -3283,6 +3417,27 @@ mod tests { use crate::query::{ExecutableQuery, QueryBase}; use crate::test_utils::connection::new_test_connection; + #[test] + fn test_tokenize_uses_explicit_simple_tokenizer() { + let params = + crate::index::scalar::FtsIndexBuilder::default().base_tokenizer("simple".to_string()); + let tokens = crate::tokenize("Running in cafés", ¶ms).unwrap(); + + assert_eq!( + tokens, + vec![ + FtsToken { + text: "run".to_string(), + position: 0, + }, + FtsToken { + text: "cafe".to_string(), + position: 2, + }, + ] + ); + } + #[tokio::test] async fn test_open() { let tmp_dir = tempdir().unwrap();