diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index f888d4b42..66aae06a6 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -591,8 +591,8 @@ class PhraseQuery(FullTextQuery): ---------- query : str The query string to match against. - column : str - The name of the column to match against. + column : str, optional + The column to match against. None leaves column selection to query planning. slop : int, default 0 The maximum number of intervening positions permitted in the phrase. document_granularity : DocumentGranularity, optional @@ -602,7 +602,7 @@ class PhraseQuery(FullTextQuery): """ query: str - column: str + column: Optional[str] = None slop: int = pydantic.Field(0, kw_only=True) document_granularity: Optional[DocumentGranularity] = pydantic.Field( None, kw_only=True diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index 1dc2fd9a1..a02c4db28 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -584,6 +584,52 @@ async def test_search_fts_phrase_query_async(async_table): assert len(phrase_results) > 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("columns", [None, ["text"]]) +async def test_phrase_query_to_query_object_async(tmp_path, columns): + db = await ldb.connect_async(tmp_path) + table = await db.create_table( + "phrases", pa.table({"text": ["puppy runs fast", "cat sleeps"]}) + ) + await table.create_index("text", config=FTS(with_position=True)) + + query = table.query().nearest_to_text('"puppy runs"', columns=columns) + phrase = query.to_query_object().full_text_query.query + + expected = ( + PhraseQuery("puppy runs", "text") if columns else PhraseQuery("puppy runs") + ) + assert phrase == expected + results = await table.query().nearest_to_text(phrase).to_list() + assert [row["text"] for row in results] == ["puppy runs fast"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("column", [None, "text"]) +async def test_phrase_query_object_roundtrip_async(tmp_path, column): + db = await ldb.connect_async(tmp_path) + table = await db.create_table( + "phrases", + pa.table({"text": ["puppy runs fast", "puppy often runs", "cat sleeps"]}), + ) + await table.create_index("text", config=FTS(with_position=True)) + phrase = PhraseQuery( + "puppy runs", column, slop=1, document_granularity=DocumentGranularity.ROW + ) + query = table.query().nearest_to_text(phrase) + + assert query.to_query_object().full_text_query.query == phrase + serialized = json.loads(phrase.to_json())["phrase"] + assert serialized["slop"] == 1 + assert serialized["document_granularity"] == "row" + assert serialized.get("column") == column + results = await query.to_list() + assert sorted(row["text"] for row in results) == [ + "puppy often runs", + "puppy runs fast", + ] + + def test_search_fts_specify_column(table): table.create_fts_index("text") table.create_fts_index("text2") diff --git a/python/src/query.rs b/python/src/query.rs index c9c4905c9..bf267443c 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -103,7 +103,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB { } "PhraseQuery" => { let query = ob.getattr("query")?.extract()?; - let column = ob.getattr("column")?.extract()?; + let column = ob.getattr("column")?.extract::>()?; let slop = ob.getattr("slop")?.extract()?; let document_granularity = ob .getattr("document_granularity")? @@ -114,9 +114,7 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB { }) .transpose()?; - let mut query = PhraseQuery::new(query) - .with_column(Some(column)) - .with_slop(slop); + let mut query = PhraseQuery::new(query).with_column(column).with_slop(slop); if let Some(document_granularity) = document_granularity { query = query.with_document_granularity(document_granularity); } @@ -213,7 +211,7 @@ impl<'py> IntoPyObject<'py> for PyLanceDB { } namespace .getattr(intern!(py, "PhraseQuery"))? - .call((query.terms, query.column.unwrap()), Some(&kwargs)) + .call((query.terms, query.column), Some(&kwargs)) } FtsQuery::Boost(query) => { let positive = Self(query.positive.as_ref().clone()).into_pyobject(py)?;