diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index 162a2f11..ff5c6287 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -167,8 +167,24 @@ def test_search_index(tmp_path, table): @pytest.mark.parametrize("use_tantivy", [True, False]) def test_search_fts(table, use_tantivy): table.create_fts_index("text", use_tantivy=use_tantivy) - results = table.search("puppy").limit(5).to_list() + results = table.search("puppy").select(["id", "text"]).limit(5).to_list() assert len(results) == 5 + assert len(results[0]) == 3 # id, text, _score + + +@pytest.mark.asyncio +async def test_fts_select_async(async_table): + tbl = await async_table + await tbl.create_index("text", config=FTS()) + results = ( + await tbl.query() + .nearest_to_text("puppy") + .select(["id", "text"]) + .limit(5) + .to_list() + ) + assert len(results) == 5 + assert len(results[0]) == 3 # id, text, _score def test_search_fts_phrase_query(table): diff --git a/python/src/query.rs b/python/src/query.rs index 2df0bf15..ff2057e9 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -152,6 +152,10 @@ impl FTSQuery { self.inner = self.inner.clone().select(Select::dynamic(&columns)); } + pub fn select_columns(&mut self, columns: Vec) { + self.inner = self.inner.clone().select(Select::columns(&columns)); + } + pub fn limit(&mut self, limit: u32) { self.inner = self.inner.clone().limit(limit as usize); } @@ -341,6 +345,11 @@ impl HybridQuery { self.inner_fts.select(columns); } + pub fn select_columns(&mut self, columns: Vec) { + self.inner_vec.select_columns(columns.clone()); + self.inner_fts.select_columns(columns); + } + pub fn limit(&mut self, limit: u32) { self.inner_vec.limit(limit); self.inner_fts.limit(limit);