From 659c64b52338c4f1e83a26a3d40a1145b40b5846 Mon Sep 17 00:00:00 2001 From: Yang Cen Date: Tue, 28 Jul 2026 22:11:58 +0800 Subject: [PATCH] feat(fts): add custom stop-word sources --- docs/src/js/globals.md | 1 + docs/src/js/interfaces/FtsOptions.md | 27 ++ .../js/type-aliases/CustomStopWordsSource.md | 34 +++ nodejs/__test__/remote.test.ts | 59 +++- nodejs/__test__/table.test.ts | 96 +++++++ nodejs/lancedb/index.ts | 1 + nodejs/lancedb/indices.ts | 62 +++- nodejs/lancedb/table.ts | 6 + nodejs/src/index.rs | 4 + nodejs/src/table.rs | 8 + python/python/lancedb/index.py | 61 +++- python/python/lancedb/remote/table.py | 5 + python/python/lancedb/table.py | 15 + python/python/tests/test_fts.py | 78 +++++- python/python/tests/test_remote_db.py | 55 +++- python/src/index.rs | 60 +++- python/src/table.rs | 5 +- rust/lancedb/src/index.rs | 22 +- rust/lancedb/src/index/scalar.rs | 30 ++ rust/lancedb/src/remote/table.rs | 264 ++++++++++++++++-- rust/lancedb/src/table.rs | 95 +++++-- rust/lancedb/src/table/create_index.rs | 180 +++++++++++- 22 files changed, 1091 insertions(+), 77 deletions(-) create mode 100644 docs/src/js/type-aliases/CustomStopWordsSource.md diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index c214a5f13..d381a12fe 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -128,6 +128,7 @@ - [AnalyzePlanDistributedMetrics](type-aliases/AnalyzePlanDistributedMetrics.md) - [BaseTokenizer](type-aliases/BaseTokenizer.md) +- [CustomStopWordsSource](type-aliases/CustomStopWordsSource.md) - [Data](type-aliases/Data.md) - [DataLike](type-aliases/DataLike.md) - [FieldLike](type-aliases/FieldLike.md) diff --git a/docs/src/js/interfaces/FtsOptions.md b/docs/src/js/interfaces/FtsOptions.md index b8b264fed..74a571fcb 100644 --- a/docs/src/js/interfaces/FtsOptions.md +++ b/docs/src/js/interfaces/FtsOptions.md @@ -56,6 +56,33 @@ the experimental FTS V3 format and may introduce breaking changes. *** +### customStopWords? + +```ts +optional customStopWords: string[]; +``` + +Custom stop words that replace the built-in language list. + +`undefined` uses the built-in list while `[]` explicitly selects no stop +words. This only affects tokenization when `removeStopWords` is true. + +*** + +### customStopWordsSource? + +```ts +optional customStopWordsSource: CustomStopWordsSource; +``` + +Request-only source resolved by a remote LanceDB service. + +This is mutually exclusive with `customStopWords`. Local native tables +reject source descriptors; load the source explicitly and pass +`customStopWords` instead. + +*** + ### language? ```ts diff --git a/docs/src/js/type-aliases/CustomStopWordsSource.md b/docs/src/js/type-aliases/CustomStopWordsSource.md new file mode 100644 index 000000000..b5562ff49 --- /dev/null +++ b/docs/src/js/type-aliases/CustomStopWordsSource.md @@ -0,0 +1,34 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / CustomStopWordsSource + +# Type Alias: CustomStopWordsSource + +```ts +type CustomStopWordsSource: object | object | object; +``` + +A request-only source for custom FTS stop words. + +Remote services resolve file and table descriptors into a stable snapshot. +Local native tables reject source descriptors. + +## Example + +```ts +const inline = { + type: "inline", + words: ["the", "a"], +} satisfies CustomStopWordsSource; +const file = { + type: "file", + uri: "s3://bucket/stop-words.txt", +} satisfies CustomStopWordsSource; +const table = { + type: "table", + table: "catalog.schema.stop_words", + column: "word", +} satisfies CustomStopWordsSource; +``` diff --git a/nodejs/__test__/remote.test.ts b/nodejs/__test__/remote.test.ts index 22562b0bc..7ed2a3eca 100644 --- a/nodejs/__test__/remote.test.ts +++ b/nodejs/__test__/remote.test.ts @@ -226,8 +226,8 @@ describe("remote connection", () => { ); }); - it("sends the FTS posting block size to remote tables", async () => { - let createIndexBody: Record | undefined; + it("serializes FTS custom stop words and sources to remote tables", async () => { + const createIndexBodies: Record[] = []; await withMockDatabase( (req, res) => { @@ -253,7 +253,7 @@ describe("remote connection", () => { raw += chunk; }); req.on("end", () => { - createIndexBody = JSON.parse(raw); + createIndexBodies.push(JSON.parse(raw)); res.writeHead(200).end(); }); return; @@ -266,12 +266,59 @@ describe("remote connection", () => { await table.createIndex("text", { config: Index.fts({ blockSize: 256 }), }); + await table.createIndex("text", { + config: Index.fts({ customStopWords: [] }), + }); + await table.createIndex("text", { + config: Index.fts({ customStopWords: ["lance", "database"] }), + }); + await table.createIndex("text", { + config: Index.fts({ + customStopWordsSource: { type: "inline", words: ["lance"] }, + }), + }); + await table.createIndex("text", { + config: Index.fts({ + customStopWordsSource: { + type: "file", + uri: "s3://bucket/stop-words.txt", + }, + }), + }); + await table.createIndex("text", { + config: Index.fts({ + customStopWordsSource: { + type: "table", + table: "stop_words", + column: "word", + }, + }), + }); }, ); - expect(createIndexBody?.["column"]).toBe("text"); - expect(createIndexBody?.["index_type"]).toBe("FTS"); - expect(createIndexBody?.["block_size"]).toBe(256); + expect(createIndexBodies[0]["column"]).toBe("text"); + expect(createIndexBodies[0]["index_type"]).toBe("FTS"); + expect(createIndexBodies[0]["block_size"]).toBe(256); + expect(createIndexBodies[0]["custom_stop_words"]).toBeNull(); + expect(createIndexBodies[1]["custom_stop_words"]).toEqual([]); + expect(createIndexBodies[2]["custom_stop_words"]).toEqual([ + "lance", + "database", + ]); + expect(createIndexBodies[3]["custom_stop_words_source"]).toEqual({ + type: "inline", + words: ["lance"], + }); + expect(createIndexBodies[4]["custom_stop_words_source"]).toEqual({ + type: "file", + uri: "s3://bucket/stop-words.txt", + }); + expect(createIndexBodies[5]["custom_stop_words_source"]).toEqual({ + type: "table", + table: "stop_words", + column: "word", + }); }); it("diffs and merges remote branches", async () => { diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index 15248f5cf..84776d574 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -2564,6 +2564,102 @@ describe.each([arrow15, arrow16, arrow17, arrow18])( ); }); + test("preserves missing, empty, and non-empty custom stop words", async () => { + const db = await connect(tmpDir.name); + + const builtIn = await db.createTable("built_in_stop_words", [ + { text: "the lance database" }, + ]); + await builtIn.createIndex("text", { + config: Index.fts({ stem: false }), + }); + await expect( + builtIn.tokenize("the lance", { column: "text" }), + ).resolves.toEqual([{ text: "lance", position: 1 }]); + + const empty = await db.createTable("empty_custom_stop_words", [ + { text: "the lance database" }, + ]); + await empty.createIndex("text", { + config: Index.fts({ stem: false, customStopWords: [] }), + }); + await expect( + empty.tokenize("the lance", { column: "text" }), + ).resolves.toEqual([ + { text: "the", position: 0 }, + { text: "lance", position: 1 }, + ]); + + const custom = await db.createTable("non_empty_custom_stop_words", [ + { text: "the lance database" }, + ]); + await custom.createIndex("text", { + config: Index.fts({ + stem: false, + customStopWords: ["lance"], + }), + }); + await expect( + custom.tokenize("the lance", { column: "text" }), + ).resolves.toEqual([{ text: "the", position: 0 }]); + await expect(custom.search("the", "fts").toArray()).resolves.toHaveLength( + 1, + ); + await expect( + custom.search("lance", "fts").toArray(), + ).resolves.toHaveLength(0); + + const disabled = await db.createTable("disabled_custom_stop_words", [ + { text: "the lance database" }, + ]); + await disabled.createIndex("text", { + config: Index.fts({ + stem: false, + removeStopWords: false, + customStopWords: ["lance"], + }), + }); + await expect( + disabled.tokenize("the lance", { column: "text" }), + ).resolves.toEqual([ + { text: "the", position: 0 }, + { text: "lance", position: 1 }, + ]); + await expect( + disabled.search("lance", "fts").toArray(), + ).resolves.toHaveLength(1); + }); + + test("rejects custom stop-word sources on local tables", async () => { + const db = await connect(tmpDir.name); + const table = await db.createTable("source_stop_words", [ + { text: "lance database" }, + ]); + + await expect( + table.createIndex("text", { + config: Index.fts({ + customStopWordsSource: { + type: "file", + uri: "file:///tmp/stop-words.txt", + }, + }), + }), + ).rejects.toThrow("only supported for remote tables"); + }); + + test("rejects both custom stop-word forms", () => { + expect(() => + Index.fts({ + customStopWords: [], + customStopWordsSource: { + type: "file", + uri: "s3://bucket/stop-words.txt", + }, + }), + ).toThrow("mutually exclusive"); + }); + test("full text search without lowercase", async () => { const db = await connect(tmpDir.name); const data = [ diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index cb91c30aa..c29f10ea2 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -118,6 +118,7 @@ export { HnswPqOptions, HnswSqOptions, FtsOptions, + CustomStopWordsSource, BaseTokenizer, } from "./indices"; diff --git a/nodejs/lancedb/indices.ts b/nodejs/lancedb/indices.ts index 9b438ef0f..6f2efec04 100644 --- a/nodejs/lancedb/indices.ts +++ b/nodejs/lancedb/indices.ts @@ -553,6 +553,23 @@ export interface FtsOptions { */ removeStopWords?: boolean; + /** + * Custom stop words that replace the built-in language list. + * + * `undefined` uses the built-in list while `[]` explicitly selects no stop + * words. This only affects tokenization when `removeStopWords` is true. + */ + customStopWords?: string[]; + + /** + * Request-only source resolved by a remote LanceDB service. + * + * This is mutually exclusive with `customStopWords`. Local native tables + * reject source descriptors; load the source explicitly and pass + * `customStopWords` instead. + */ + customStopWordsSource?: CustomStopWordsSource; + /** * whether to remove punctuation */ @@ -582,10 +599,43 @@ export interface FtsOptions { blockSize?: 128 | 256; } +/** + * A request-only source for custom FTS stop words. + * + * Remote services resolve file and table descriptors into a stable snapshot. + * Local native tables reject source descriptors. + * + * @example + * ```ts + * const inline = { + * type: "inline", + * words: ["the", "a"], + * } satisfies CustomStopWordsSource; + * const file = { + * type: "file", + * uri: "s3://bucket/stop-words.txt", + * } satisfies CustomStopWordsSource; + * const table = { + * type: "table", + * table: "catalog.schema.stop_words", + * column: "word", + * } satisfies CustomStopWordsSource; + * ``` + */ +export type CustomStopWordsSource = + | { type: "inline"; words: string[] } + | { type: "file"; uri: string } + | { type: "table"; table: string; column: string }; + export class Index { private readonly inner: LanceDbIndex; - private constructor(inner: LanceDbIndex) { + private readonly customStopWordsSource?: CustomStopWordsSource; + private constructor( + inner: LanceDbIndex, + customStopWordsSource?: CustomStopWordsSource, + ) { this.inner = inner; + this.customStopWordsSource = customStopWordsSource; } /** @@ -746,6 +796,14 @@ export class Index { * You can combine filters with full text search. */ static fts(options?: Partial) { + if ( + options?.customStopWords !== undefined && + options?.customStopWordsSource !== undefined + ) { + throw new Error( + "customStopWords and customStopWordsSource are mutually exclusive; choose one", + ); + } return new Index( LanceDbIndex.fts( options?.withPosition, @@ -760,7 +818,9 @@ export class Index { options?.ngramMaxLength, options?.prefixOnly, options?.blockSize, + options?.customStopWords, ), + options?.customStopWordsSource, ); } diff --git a/nodejs/lancedb/table.ts b/nodejs/lancedb/table.ts index d814e3e55..be93bd318 100644 --- a/nodejs/lancedb/table.ts +++ b/nodejs/lancedb/table.ts @@ -930,6 +930,9 @@ export class LocalTable extends Table { // Bit of a hack to get around the fact that TS has no package-scope. // biome-ignore lint/suspicious/noExplicitAny: skip const nativeIndex = (options?.config as any)?.inner; + // biome-ignore lint/suspicious/noExplicitAny: skip + const customStopWordsSource = (options?.config as any) + ?.customStopWordsSource; await this.inner.createIndex( nativeIndex, column, @@ -937,6 +940,9 @@ export class LocalTable extends Table { options?.waitTimeoutSeconds, options?.name, options?.train, + customStopWordsSource === undefined + ? undefined + : JSON.stringify(customStopWordsSource), ); } diff --git a/nodejs/src/index.rs b/nodejs/src/index.rs index 66dee7913..d115a37b5 100644 --- a/nodejs/src/index.rs +++ b/nodejs/src/index.rs @@ -227,6 +227,7 @@ impl Index { ngram_max_length: Option, prefix_only: Option, block_size: Option, + custom_stop_words: Option>, ) -> napi::Result { let mut opts = FtsIndexBuilder::default(); if let Some(with_position) = with_position { @@ -267,6 +268,9 @@ impl Index { .block_size(block_size as usize) .map_err(|err| napi::Error::from_reason(err.to_string()))?; } + if let Some(custom_stop_words) = custom_stop_words { + opts = opts.custom_stop_words(Some(custom_stop_words)); + } Ok(Self { inner: Mutex::new(Some(LanceDbIndex::FTS(opts))), diff --git a/nodejs/src/table.rs b/nodejs/src/table.rs index d1489034f..388b1ae61 100644 --- a/nodejs/src/table.rs +++ b/nodejs/src/table.rs @@ -137,6 +137,7 @@ impl Table { } #[napi(catch_unwind)] + #[allow(clippy::too_many_arguments)] pub async fn create_index( &self, index: Option<&Index>, @@ -145,6 +146,7 @@ impl Table { wait_timeout_s: Option, name: Option, train: Option, + custom_stop_words_source: Option, ) -> napi::Result<()> { let lancedb_index = if let Some(index) = index { index.consume()? @@ -165,6 +167,12 @@ impl Table { if let Some(train) = train { builder = builder.train(train); } + if let Some(source) = custom_stop_words_source { + let source = serde_json::from_str(&source).map_err(|err| { + napi::Error::from_reason(format!("invalid customStopWordsSource: {err}")) + })?; + builder = builder.custom_stop_words_source(source); + } builder.execute().await.default_error() } diff --git a/python/python/lancedb/index.py b/python/python/lancedb/index.py index c06bb9903..49d330878 100644 --- a/python/python/lancedb/index.py +++ b/python/python/lancedb/index.py @@ -1,8 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The LanceDB Authors -from dataclasses import dataclass -from typing import Literal, Optional +from dataclasses import dataclass, field +from typing import List, Literal, Optional, Union from ._lancedb import ( IndexConfig, @@ -31,6 +31,38 @@ lang_mapping = { } +@dataclass(frozen=True) +class InlineStopWordsSource: + """Inline custom stop words for a remote FTS create-index request.""" + + words: List[str] + type: Literal["inline"] = field(default="inline", init=False) + + +@dataclass(frozen=True) +class FileStopWordsSource: + """A UTF-8, newline-delimited custom stop-word object for a remote request.""" + + uri: str + type: Literal["file"] = field(default="file", init=False) + + +@dataclass(frozen=True) +class TableStopWordsSource: + """A table string column used by a remote FTS create-index request.""" + + table: str + column: str + type: Literal["table"] = field(default="table", init=False) + + +CustomStopWordsSource = Union[ + InlineStopWordsSource, + FileStopWordsSource, + TableStopWordsSource, +] + + @dataclass class BTree: """Describes a btree index configuration @@ -151,6 +183,15 @@ class FTS: remove_stop_words : bool, default True Whether to remove stop words. Stop words are common words that are often removed from text before indexing. For example, in English "the" and "and". + custom_stop_words : list of str, optional + A custom stop-word list that replaces the built-in language list. ``None`` + uses the built-in list, while ``[]`` explicitly selects no stop words. + This only affects tokenization when ``remove_stop_words`` is True. + custom_stop_words_source : CustomStopWordsSource, optional + A request-only inline, UTF-8 file, or table-column source resolved by a + remote LanceDB service. It is mutually exclusive with + ``custom_stop_words``. Local native tables reject source descriptors; + load the source explicitly and pass ``custom_stop_words`` instead. ascii_folding : bool, default True Whether to fold ASCII characters. This converts accented characters to their ASCII equivalent. For example, "café" would be converted to "cafe". @@ -179,6 +220,18 @@ class FTS: ngram_max_length: int = 3 prefix_only: bool = False block_size: int = 128 + custom_stop_words: Optional[List[str]] = None + custom_stop_words_source: Optional[CustomStopWordsSource] = None + + def __post_init__(self): + if ( + self.custom_stop_words is not None + and self.custom_stop_words_source is not None + ): + raise ValueError( + "custom_stop_words and custom_stop_words_source are mutually " + "exclusive; choose one" + ) @dataclass @@ -853,6 +906,10 @@ __all__ = [ "HnswFlat", "IndexConfig", "FTS", + "CustomStopWordsSource", + "InlineStopWordsSource", + "FileStopWordsSource", + "TableStopWordsSource", "Bitmap", "LabelList", "Fm", diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 682b8533c..d0f6f1127 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -39,6 +39,7 @@ from lancedb.index import ( FTS, BTree, Bitmap, + CustomStopWordsSource, HnswFlat, HnswSq, IvfFlat, @@ -340,6 +341,8 @@ class RemoteTable(Table): lower_case: bool = True, stem: bool = True, remove_stop_words: bool = True, + custom_stop_words: Optional[List[str]] = None, + custom_stop_words_source: Optional[CustomStopWordsSource] = None, ascii_folding: bool = True, ngram_min_length: int = 3, ngram_max_length: int = 3, @@ -361,6 +364,8 @@ class RemoteTable(Table): lower_case=lower_case, stem=stem, remove_stop_words=remove_stop_words, + custom_stop_words=custom_stop_words, + custom_stop_words_source=custom_stop_words_source, ascii_folding=ascii_folding, ngram_min_length=ngram_min_length, ngram_max_length=ngram_max_length, diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 6c3854734..845d431ac 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -69,6 +69,7 @@ from .index import ( HnswSq, HnswFlat, FTS, + CustomStopWordsSource, ) from .expr import Expr from .merge import LanceMergeInsertBuilder @@ -1103,6 +1104,8 @@ class Table(ABC): lower_case: bool = True, stem: bool = True, remove_stop_words: bool = True, + custom_stop_words: Optional[List[str]] = None, + custom_stop_words_source: Optional[CustomStopWordsSource] = None, ascii_folding: bool = True, ngram_min_length: int = 3, ngram_max_length: int = 3, @@ -1170,6 +1173,12 @@ class Table(ABC): remove_stop_words : bool, default True Whether to remove stop words. Stop words are common words that are often removed from text before indexing. For example, in English "the" and "and". + custom_stop_words : list of str, optional + Custom words that replace the built-in language stop-word list. ``None`` + uses the built-in list; ``[]`` selects no stop words. + custom_stop_words_source : CustomStopWordsSource, optional + Request-only remote source for custom stop words. Mutually exclusive + with ``custom_stop_words`` and rejected by local native tables. ascii_folding : bool, default True Whether to fold ASCII characters. This converts accented characters to their ASCII equivalent. For example, "café" would be converted to "cafe". @@ -3055,6 +3064,8 @@ class LanceTable(Table): lower_case: bool = True, stem: bool = True, remove_stop_words: bool = True, + custom_stop_words: Optional[List[str]] = None, + custom_stop_words_source: Optional[CustomStopWordsSource] = None, ascii_folding: bool = True, ngram_min_length: int = 3, ngram_max_length: int = 3, @@ -3101,6 +3112,8 @@ class LanceTable(Table): "lower_case": lower_case, "stem": stem, "remove_stop_words": remove_stop_words, + "custom_stop_words": custom_stop_words, + "custom_stop_words_source": custom_stop_words_source, "ascii_folding": ascii_folding, "ngram_min_length": ngram_min_length, "ngram_max_length": ngram_max_length, @@ -3108,6 +3121,8 @@ class LanceTable(Table): } else: tokenizer_configs = self.infer_tokenizer_configs(tokenizer_name) + tokenizer_configs["custom_stop_words"] = custom_stop_words + tokenizer_configs["custom_stop_words_source"] = custom_stop_words_source config = FTS(block_size=block_size, **tokenizer_configs) diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index 127510bc2..918008369 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -22,7 +22,7 @@ import zipfile import lancedb as ldb from lancedb.db import DBConnection -from lancedb.index import FTS +from lancedb.index import FTS, FileStopWordsSource from lancedb.query import ( BoostQuery, MatchQuery, @@ -243,6 +243,82 @@ def test_create_inverted_index_rejects_invalid_block_size(table): table.create_index("text", config=FTS(block_size=129)) +def test_custom_stop_words_preserve_none_and_empty(mem_db: DBConnection): + built_in = mem_db.create_table( + "built_in_stop_words", [{"text": "the lance database"}] + ) + built_in.create_index( + "text", + config=FTS(stem=False, custom_stop_words=None), + ) + assert [token.text for token in built_in.tokenize("the lance", column="text")] == [ + "lance" + ] + + empty = mem_db.create_table( + "empty_custom_stop_words", [{"text": "the lance database"}] + ) + empty.create_index( + "text", + config=FTS(stem=False, custom_stop_words=[]), + ) + assert [token.text for token in empty.tokenize("the lance", column="text")] == [ + "the", + "lance", + ] + + custom = mem_db.create_table( + "non_empty_custom_stop_words", [{"text": "the lance database"}] + ) + custom.create_index( + "text", + config=FTS(stem=False, custom_stop_words=["lance"]), + ) + assert [token.text for token in custom.tokenize("the lance", column="text")] == [ + "the" + ] + assert len(custom.search("the", query_type="fts").to_list()) == 1 + assert custom.search("lance", query_type="fts").to_list() == [] + + disabled = mem_db.create_table( + "disabled_custom_stop_words", [{"text": "the lance database"}] + ) + disabled.create_index( + "text", + config=FTS( + stem=False, + remove_stop_words=False, + custom_stop_words=["lance"], + ), + ) + assert [token.text for token in disabled.tokenize("the lance", column="text")] == [ + "the", + "lance", + ] + assert len(disabled.search("lance", query_type="fts").to_list()) == 1 + + +def test_local_fts_rejects_custom_stop_words_source(mem_db: DBConnection): + table = mem_db.create_table("source_stop_words", [{"text": "lance database"}]) + with pytest.raises(RuntimeError, match="only supported for remote tables"): + table.create_index( + "text", + config=FTS( + custom_stop_words_source=FileStopWordsSource( + "file:///tmp/stop-words.txt" + ) + ), + ) + + +def test_fts_rejects_both_custom_stop_words_forms(): + with pytest.raises(ValueError, match="mutually exclusive"): + FTS( + custom_stop_words=[], + custom_stop_words_source=FileStopWordsSource("s3://bucket/stop-words.txt"), + ) + + def test_search_fts(table): table.create_fts_index("text") results = table.search("puppy").select(["id", "text"]).limit(5).to_list() diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 648a52b7a..1418ad670 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -868,7 +868,15 @@ def test_remote_create_index_new_api(): request.send_response(404) request.end_headers() - from lancedb.index import BTree, FTS, IvfPq, IvfRq + from lancedb.index import ( + BTree, + FTS, + FileStopWordsSource, + InlineStopWordsSource, + IvfPq, + IvfRq, + TableStopWordsSource, + ) with mock_lancedb_connection(handler) as db: table = db.create_table("test", [{"id": 1}]) @@ -881,6 +889,28 @@ def test_remote_create_index_new_api(): table.create_index("vector", config=IvfPq(distance_type="l2")) table.create_index("category", config=BTree()) table.create_index("text", config=FTS(block_size=256)) + table.create_index("text", config=FTS(custom_stop_words=[])) + table.create_index( + "text", config=FTS(custom_stop_words=["lance", "database"]) + ) + table.create_index( + "text", + config=FTS(custom_stop_words_source=InlineStopWordsSource(["lance"])), + ) + table.create_index( + "text", + config=FTS( + custom_stop_words_source=FileStopWordsSource( + "s3://bucket/stop-words.txt" + ) + ), + ) + table.create_index( + "text", + config=FTS( + custom_stop_words_source=TableStopWordsSource("stop_words", "word") + ), + ) # IvfRq via new API table.create_index("vector", config=IvfRq(distance_type="l2")) @@ -892,15 +922,36 @@ def test_remote_create_index_new_api(): num_partitions=8, ) - assert len(received_requests) == 5 + assert len(received_requests) == 10 assert [req["column"] for req in received_requests] == [ "vector", "category", "text", + "text", + "text", + "text", + "text", + "text", "vector", "vector", ] assert received_requests[2]["block_size"] == 256 + assert received_requests[2]["custom_stop_words"] is None + assert received_requests[3]["custom_stop_words"] == [] + assert received_requests[4]["custom_stop_words"] == ["lance", "database"] + assert received_requests[5]["custom_stop_words_source"] == { + "type": "inline", + "words": ["lance"], + } + assert received_requests[6]["custom_stop_words_source"] == { + "type": "file", + "uri": "s3://bucket/stop-words.txt", + } + assert received_requests[7]["custom_stop_words_source"] == { + "type": "table", + "table": "stop_words", + "column": "word", + } def test_table_wait_for_index_timeout(): diff --git a/python/src/index.rs b/python/src/index.rs index a98b90cf9..dac24f463 100644 --- a/python/src/index.rs +++ b/python/src/index.rs @@ -8,7 +8,7 @@ use lancedb::index::vector::{ }; use lancedb::index::{ Index as LanceDbIndex, - scalar::{BTreeIndexBuilder, FmIndexBuilder, FtsIndexBuilder}, + scalar::{BTreeIndexBuilder, CustomStopWordsSource, FmIndexBuilder, FtsIndexBuilder}, }; use pyo3::IntoPyObject; use pyo3::types::PyStringMethods; @@ -33,14 +33,42 @@ pub fn class_name(ob: &'_ Bound<'_, PyAny>) -> PyResult { } } -pub fn extract_index_params(source: &Option>) -> PyResult { +fn extract_custom_stop_words_source( + source: &Bound<'_, PyAny>, +) -> PyResult> { + let source = source.getattr("custom_stop_words_source")?; + if source.is_none() { + return Ok(None); + } + let source_type: String = source.getattr("type")?.extract()?; + match source_type.as_str() { + "inline" => Ok(Some(CustomStopWordsSource::Inline { + words: source.getattr("words")?.extract()?, + })), + "file" => Ok(Some(CustomStopWordsSource::File { + uri: source.getattr("uri")?.extract()?, + })), + "table" => Ok(Some(CustomStopWordsSource::Table { + table: source.getattr("table")?.extract()?, + column: source.getattr("column")?.extract()?, + })), + other => Err(PyValueError::new_err(format!( + "invalid custom_stop_words_source type '{other}'; expected inline, file, or table" + ))), + } +} + +pub fn extract_index_params( + source: &Option>, +) -> PyResult<(LanceDbIndex, Option)> { if let Some(source) = source { match class_name(source)?.as_str() { - "BTree" => Ok(LanceDbIndex::BTree(BTreeIndexBuilder::default())), - "Bitmap" => Ok(LanceDbIndex::Bitmap(Default::default())), - "LabelList" => Ok(LanceDbIndex::LabelList(Default::default())), - "Fm" => Ok(LanceDbIndex::Fm(FmIndexBuilder::default())), + "BTree" => Ok((LanceDbIndex::BTree(BTreeIndexBuilder::default()), None)), + "Bitmap" => Ok((LanceDbIndex::Bitmap(Default::default()), None)), + "LabelList" => Ok((LanceDbIndex::LabelList(Default::default()), None)), + "Fm" => Ok((LanceDbIndex::Fm(FmIndexBuilder::default()), None)), "FTS" => { + let custom_stop_words_source = extract_custom_stop_words_source(source)?; let params = source.extract::()?; let inner_opts = FtsIndexBuilder::default() .base_tokenizer(params.base_tokenizer) @@ -55,6 +83,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult>) -> PyResult { let params = source.extract::()?; @@ -79,7 +108,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult { let params = source.extract::()?; @@ -98,7 +127,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult { let params = source.extract::()?; @@ -113,7 +142,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult { let params = source.extract::()?; @@ -129,7 +158,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult { let params = source.extract::()?; @@ -150,7 +179,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult { let params = source.extract::()?; @@ -167,7 +196,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult { let params = source.extract::()?; @@ -185,7 +214,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult Err(PyValueError::new_err(format!( "Invalid index type '{}'. Must be one of BTree, Bitmap, LabelList, Fm, FTS, IvfPq, IvfSq, IvfHnswPq, IvfHnswSq, or IvfHnswFlat", @@ -193,7 +222,7 @@ pub fn extract_index_params(source: &Option>) -> PyResult>, ascii_folding: bool, ngram_min_length: u32, ngram_max_length: u32, diff --git a/python/src/table.rs b/python/src/table.rs index bb44fc023..aaeae662b 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -781,7 +781,7 @@ impl Table { name: Option, train: Option, ) -> PyResult> { - let index = extract_index_params(&index)?; + let (index, custom_stop_words_source) = extract_index_params(&index)?; let timeout = wait_timeout.map(|t| t.extract::().unwrap()); let mut op = self_ .inner_ref()? @@ -795,6 +795,9 @@ impl Table { if let Some(train) = train { op = op.train(train); } + if let Some(source) = custom_stop_words_source { + op = op.custom_stop_words_source(source); + } future_into_py(self_.py(), async move { op.execute().await.infer_error()?; diff --git a/rust/lancedb/src/index.rs b/rust/lancedb/src/index.rs index 308a64ca1..7e815f44d 100644 --- a/rust/lancedb/src/index.rs +++ b/rust/lancedb/src/index.rs @@ -13,7 +13,10 @@ use crate::index::vector::IvfRqIndexBuilder; use crate::{DistanceType, Error, Result, table::BaseTable}; use self::{ - scalar::{BTreeIndexBuilder, BitmapIndexBuilder, FmIndexBuilder, LabelListIndexBuilder}, + scalar::{ + BTreeIndexBuilder, BitmapIndexBuilder, CustomStopWordsSource, FmIndexBuilder, + LabelListIndexBuilder, + }, vector::{ IvfHnswFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, IvfPqIndexBuilder, IvfSqIndexBuilder, @@ -192,6 +195,7 @@ pub struct IndexBuilder { pub(crate) wait_timeout: Option, pub(crate) train: bool, pub(crate) name: Option, + pub(crate) custom_stop_words_source: Option, } impl IndexBuilder { @@ -204,6 +208,7 @@ impl IndexBuilder { train: true, wait_timeout: None, name: None, + custom_stop_words_source: None, } } @@ -293,6 +298,21 @@ impl IndexBuilder { self } + /// Supply a request-only source for custom FTS stop words. + /// + /// This option is supported only for FTS indexes on remote tables. The + /// remote service resolves the source to a stable snapshot before the + /// distributed build starts. Local native tables return an error instead of + /// reading a path or table implicitly; load the source explicitly and use + /// [`FtsIndexBuilder::custom_stop_words`] for local indexes. + /// + /// This is mutually exclusive with inline custom stop words configured + /// directly on the [`FtsIndexBuilder`]. + pub fn custom_stop_words_source(mut self, source: CustomStopWordsSource) -> Self { + self.custom_stop_words_source = Some(source); + self + } + /// Duration of time to wait for asynchronous indexing to complete. If not set, /// `create_index()` will not wait. /// diff --git a/rust/lancedb/src/index/scalar.rs b/rust/lancedb/src/index/scalar.rs index 10d835bb1..87437a396 100644 --- a/rust/lancedb/src/index/scalar.rs +++ b/rust/lancedb/src/index/scalar.rs @@ -10,6 +10,36 @@ //! etc. Scalar indices can also speed up prefiltering for vector searches. A single //! vector search with prefiltering can use both a scalar index and a vector index. +use serde::{Deserialize, Serialize}; + +/// A request-only source for custom FTS stop words. +/// +/// Remote LanceDB services resolve this source to a stable inline snapshot before +/// starting an index build. Local native tables do not resolve external sources; +/// callers must load them themselves and use +/// [`FtsIndexBuilder::custom_stop_words`] instead. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case", deny_unknown_fields)] +pub enum CustomStopWordsSource { + /// Stop words supplied directly in the create-index request. + Inline { + /// Words that replace the built-in language stop-word list. + words: Vec, + }, + /// A strict UTF-8, newline-delimited object. + File { + /// URI of the object to read. + uri: String, + }, + /// A string column in a table visible to the remote service. + Table { + /// Logical table identifier. + table: String, + /// Top-level string column containing one stop word per row. + column: String, + }, +} + /// Builder for a btree index /// /// A btree index is an index on scalar columns. The index stores a copy of the column diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 4c3c1bc6f..f461ca70c 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -2318,6 +2318,27 @@ impl BaseTable for RemoteTable { } }; + if let Some(source) = index.custom_stop_words_source.as_ref() { + if !matches!(&index.index, Index::FTS(_)) { + return Err(Error::InvalidInput { + message: "`custom_stop_words_source` can only be used with an FTS index" + .to_string(), + }); + } + if params + .as_ref() + .and_then(|params| params.get("custom_stop_words")) + .is_some_and(|value| !value.is_null()) + { + return Err(Error::InvalidInput { + message: "`custom_stop_words` and `custom_stop_words_source` are mutually \ + exclusive; choose one" + .to_string(), + }); + } + body["custom_stop_words_source"] = to_json(source)?; + } + body[INDEX_TYPE_KEY] = index_type_str.into(); if let Some(params) = params { for (key, value) in params.as_object().expect("params should be a JSON object") { @@ -2956,7 +2977,7 @@ mod tests { use super::*; use crate::remote::client::{ClientConfig, RetryConfig}; - use crate::table::{AddDataMode, FieldMetadataUpdate, FtsToken}; + use crate::table::{AddDataMode, FieldMetadataUpdate}; use arrow::{array::AsArray, compute::concat_batches, datatypes::Int32Type}; use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, record_batch}; @@ -2969,6 +2990,7 @@ mod tests { use rstest::rstest; use serde_json::json; + use crate::index::scalar::CustomStopWordsSource; use crate::index::vector::{ IvfFlatIndexBuilder, IvfHnswFlatIndexBuilder, IvfHnswSqIndexBuilder, IvfRqIndexBuilder, IvfSqIndexBuilder, @@ -4553,6 +4575,14 @@ mod tests { }, Index::FTS(InvertedIndexParams::default().block_size(256).unwrap()), ), + ( + "FTS", + serde_json::to_value( + InvertedIndexParams::default().custom_stop_words(Some(Vec::new())), + ) + .unwrap(), + Index::FTS(InvertedIndexParams::default().custom_stop_words(Some(Vec::new()))), + ), ]; for (index_type, expected_body, index) in cases { @@ -4592,6 +4622,137 @@ mod tests { } } + #[tokio::test] + async fn test_create_fts_index_with_custom_stop_words_source() { + let sources = [ + ( + CustomStopWordsSource::Inline { + words: vec!["lance".to_string()], + }, + json!({"type": "inline", "words": ["lance"]}), + ), + ( + CustomStopWordsSource::File { + uri: "s3://bucket/stop-words.txt".to_string(), + }, + json!({"type": "file", "uri": "s3://bucket/stop-words.txt"}), + ), + ( + CustomStopWordsSource::Table { + table: "stop_words".to_string(), + column: "word".to_string(), + }, + json!({"type": "table", "table": "stop_words", "column": "word"}), + ), + ]; + + for (source, expected_source) in sources { + let table = + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["custom_stop_words_source"], expected_source); + assert!(body["custom_stop_words"].is_null()); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + }); + + table + .create_index(&["text"], Index::FTS(InvertedIndexParams::default())) + .custom_stop_words_source(source) + .execute() + .await + .unwrap(); + } + } + + #[tokio::test] + async fn test_create_fts_index_custom_stop_words_serialization() { + let cases = [ + (None, serde_json::Value::Null), + (Some(Vec::new()), json!([])), + ( + Some(vec!["lance".to_string(), "database".to_string()]), + json!(["lance", "database"]), + ), + ]; + + for (custom_stop_words, expected) in cases { + let table = + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + "/v1/table/my_table/create_index/" => { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["custom_stop_words"], expected); + http::Response::builder() + .status(200) + .body("{}".to_string()) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + }); + + table + .create_index( + &["text"], + Index::FTS(InvertedIndexParams::default().custom_stop_words(custom_stop_words)), + ) + .execute() + .await + .unwrap(); + } + } + + #[tokio::test] + async fn test_create_fts_index_rejects_both_custom_stop_words_forms() { + let table = + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + http::Response::builder() + .status(200) + .body(describe_response(&schema)) + .unwrap() + } + path => panic!("Unexpected path: {path}"), + }); + + let error = table + .create_index( + &["text"], + Index::FTS( + InvertedIndexParams::default() + .custom_stop_words(Some(vec!["lance".to_string()])), + ), + ) + .custom_stop_words_source(CustomStopWordsSource::Inline { words: Vec::new() }) + .execute() + .await + .unwrap_err(); + assert!(error.to_string().contains("mutually exclusive")); + } + #[tokio::test] async fn test_create_index_nested_field_paths() { let schema = nested_index_schema(); @@ -5075,17 +5236,73 @@ mod tests { } #[tokio::test] - async fn test_tokenize_uses_remote_index_details() { + async fn test_tokenize_preserves_remote_custom_stop_words_tristate() { + let cases = [ + (serde_json::Value::Null, vec!["lance"]), + (serde_json::json!([]), vec!["the", "lance"]), + (serde_json::json!(["lance"]), vec!["the"]), + ]; + + for (custom_stop_words, expected) in cases { + let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); + let index_details = serde_json::json!({ + "lance_tokenizer": null, + "base_tokenizer": "simple", + "language": "English", + "with_position": false, + "max_token_length": 40, + "lower_case": true, + "stem": false, + "remove_stop_words": true, + "ascii_folding": true, + "custom_stop_words": custom_stop_words, + }) + .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("the lance", "text_idx").await.unwrap(); + assert_eq!( + tokens + .iter() + .map(|token| token.text.as_str()) + .collect::>(), + expected + ); + } + } + + #[tokio::test] + async fn test_tokenize_rejects_lossy_remote_index_details() { let schema = Schema::new(vec![Field::new("text", DataType::Utf8, false)]); let index_details = serde_json::json!({ - "base_tokenizer": "icu", + "base_tokenizer": "simple", "language": "English", - "with_position": false, - "max_token_length": 40, - "lower_case": true, - "stem": false, - "remove_stop_words": false, - "ascii_folding": true, + "remove_stop_words": true, }) .to_string(); let table = Table::new_with_handler("my_table", move |request| { @@ -5115,28 +5332,13 @@ mod tests { } }); - 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, - }, - ] - ); + let err = table.tokenize("the lance", "text_idx").await.unwrap_err(); + assert!(matches!( + err, + Error::InvalidInput { message } + if message.contains("incomplete remote tokenizer details") + && message.contains("custom_stop_words") + )); } #[tokio::test] diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 0e9354ffc..19734d9d7 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -1765,7 +1765,10 @@ impl Table { /// /// 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. + /// same tokenizer model files must also exist locally. Remote index details + /// must contain the full persisted tokenizer parameters, including an + /// explicit `custom_stop_words` field. Tokenization fails closed when the + /// server only returns lossy index statistics. pub async fn tokenize(&self, query: &str, index_name: &str) -> Result> { let indices = self.inner.list_indices().await?; let matches = indices @@ -1790,7 +1793,7 @@ impl Table { message: format!("Index '{}' is not a full text search index", index_name), }); } - self.tokenize_with_index(query, index, index_name) + self.tokenize_with_index(query, index, index_name).await } /// Tokenize a full-text search query using the tokenizer configured on the @@ -1799,7 +1802,9 @@ impl Table { /// 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. + /// files must also exist locally. Remote index details must contain the + /// full persisted tokenizer parameters, including an explicit + /// `custom_stop_words` field. 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)?; @@ -1831,30 +1836,80 @@ impl Table { self.tokenize(query, &index.name).await } - fn tokenize_with_index( + async 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 { + let params = if let Some(native) = self.inner.as_native() { + let dataset = native.dataset.get().await?; + let segments = dataset.load_indices_by_name(index_name).await?; + let mut params = None; + for segment in segments.iter() { + let segment_params = + lance::index::scalar::load_segment_params(&dataset, segment).await?; + match ¶ms { + Some(expected) if expected != &segment_params => { + return Err(Error::InvalidInput { + message: format!( + "Full text search index '{}' for {} has inconsistent tokenizer \ + parameters across segments", + index.name, selector_description + ), + }); + } + Some(_) => {} + None => params = Some(segment_params), + } + } + params.ok_or_else(|| Error::InvalidInput { message: format!( - "Full text search index '{}' for {} does not include tokenizer details", + "Full text search index '{}' for {} does not have any committed segments", index.name, selector_description ), + })? + } else { + 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 details = 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 + ), + } })?; - 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 - ), + if !details.as_object().is_some_and(|details| { + details.contains_key("lance_tokenizer") && details.contains_key("custom_stop_words") + }) { + return Err(Error::InvalidInput { + message: format!( + "Full text search index '{}' for {} has incomplete remote tokenizer \ + details: the server must return the full persisted parameters \ + (including `lance_tokenizer` and `custom_stop_words`); tokenization \ + cannot safely fall back to language defaults", + index.name, selector_description + ), + }); } - })?; + serde_json::from_value::(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!( @@ -3008,6 +3063,14 @@ impl BaseTable for NativeTable { } async fn create_index(&self, opts: IndexBuilder) -> Result<()> { + if opts.custom_stop_words_source.is_some() { + return Err(Error::NotSupported { + message: "`custom_stop_words_source` is only supported for remote tables; \ + resolve file or table sources to `custom_stop_words` before creating \ + a local native FTS index" + .to_string(), + }); + } if opts.columns.len() != 1 { return Err(Error::Schema { message: "Multi-column (composite) indices are not yet supported".to_string(), diff --git a/rust/lancedb/src/table/create_index.rs b/rust/lancedb/src/table/create_index.rs index 6be9de0e6..08e5eff63 100644 --- a/rust/lancedb/src/table/create_index.rs +++ b/rust/lancedb/src/table/create_index.rs @@ -383,7 +383,8 @@ mod tests { use crate::connection::ConnectBuilder; use crate::index::Index; use crate::index::scalar::{ - BTreeIndexBuilder, BitmapIndexBuilder, FmIndexBuilder, FtsIndexBuilder, + BTreeIndexBuilder, BitmapIndexBuilder, CustomStopWordsSource, FmIndexBuilder, + FtsIndexBuilder, }; use crate::index::vector::{ IvfHnswFlatIndexBuilder, IvfHnswPqIndexBuilder, IvfHnswSqIndexBuilder, @@ -1393,6 +1394,183 @@ mod tests { table.prewarm_index("text_idx").await.unwrap(); } + #[tokio::test] + async fn test_native_fts_rejects_custom_stop_words_source() { + let conn = connect("memory://").execute().await.unwrap(); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("text", DataType::Utf8, false)])), + vec![Arc::new(StringArray::from(vec!["lance database"])) as ArrayRef], + ) + .unwrap(); + let table = conn + .create_table("test_custom_stop_words_source", batch) + .execute() + .await + .unwrap(); + + let error = table + .create_index(&["text"], Index::FTS(FtsIndexBuilder::default())) + .custom_stop_words_source(CustomStopWordsSource::File { + uri: "file:///tmp/stop-words.txt".to_string(), + }) + .execute() + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("only supported for remote tables") + ); + assert!(error.to_string().contains("custom_stop_words")); + } + + #[tokio::test] + async fn test_fts_custom_stop_words_survive_table_reopen() { + let tempdir = tempdir().unwrap(); + let uri = tempdir.path().to_str().unwrap(); + let conn = connect(uri).execute().await.unwrap(); + + for (name, custom_stop_words) in [ + ("default_stop_words", None), + ("empty_custom_stop_words", Some(Vec::::new())), + ( + "non_empty_custom_stop_words", + Some(vec!["lance".to_string()]), + ), + ] { + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("text", DataType::Utf8, false)])), + vec![Arc::new(StringArray::from(vec!["the lance database"])) as ArrayRef], + ) + .unwrap(); + let table = conn.create_table(name, batch).execute().await.unwrap(); + table + .create_index( + &["text"], + Index::FTS( + FtsIndexBuilder::default() + .stem(false) + .custom_stop_words(custom_stop_words), + ), + ) + .execute() + .await + .unwrap(); + } + drop(conn); + + let reopened = connect(uri).execute().await.unwrap(); + let default = reopened + .open_table("default_stop_words") + .execute() + .await + .unwrap(); + let default_tokens = default + .tokenize_with_column("the lance", "text") + .await + .unwrap(); + assert_eq!( + default_tokens + .iter() + .map(|token| token.text.as_str()) + .collect::>(), + vec!["lance"] + ); + let default_results = default + .query() + .full_text_search(FullTextSearchQuery::new("the".to_string())) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + default_results + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 0 + ); + + let empty = reopened + .open_table("empty_custom_stop_words") + .execute() + .await + .unwrap(); + let empty_tokens = empty + .tokenize_with_column("the lance", "text") + .await + .unwrap(); + assert_eq!( + empty_tokens + .iter() + .map(|token| token.text.as_str()) + .collect::>(), + vec!["the", "lance"] + ); + let empty_results = empty + .query() + .full_text_search(FullTextSearchQuery::new("the".to_string())) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + empty_results + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 1 + ); + + let custom = reopened + .open_table("non_empty_custom_stop_words") + .execute() + .await + .unwrap(); + let custom_tokens = custom + .tokenize_with_column("the lance", "text") + .await + .unwrap(); + assert_eq!( + custom_tokens + .iter() + .map(|token| token.text.as_str()) + .collect::>(), + vec!["the"] + ); + + let results = custom + .query() + .full_text_search(FullTextSearchQuery::new("the".to_string())) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(results.iter().map(RecordBatch::num_rows).sum::(), 1); + let removed_results = custom + .query() + .full_text_search(FullTextSearchQuery::new("lance".to_string())) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!( + removed_results + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 0 + ); + } + #[tokio::test] pub async fn test_list_indices_skip_frag_reuse() { let conn = connect("memory://").execute().await.unwrap();