mirror of
https://github.com/lancedb/lancedb.git
synced 2026-01-08 12:52:58 +00:00
fix: bugs for new FTS APIs (#2314)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced full-text search capabilities with support for phrase queries, fuzzy matching, boosting, and multi-column matching. - Search methods now accept full-text query objects directly, improving query flexibility and precision. - Python and JavaScript SDKs updated to handle full-text queries seamlessly, including async search support. - **Tests** - Added comprehensive tests covering fuzzy search, phrase search, and boosted queries to ensure robust full-text search functionality. - **Documentation** - Updated query class documentation to reflect new constructor options and removal of deprecated methods for clarity and simplicity. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: BubbleCal <bubble-cal@outlook.com>
This commit is contained in:
@@ -3,7 +3,9 @@
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use lancedb::index::scalar::{FtsQuery, FullTextSearchQuery, MatchQuery, PhraseQuery};
|
||||
use lancedb::index::scalar::{
|
||||
BoostQuery, FtsQuery, FullTextSearchQuery, MatchQuery, MultiMatchQuery, PhraseQuery,
|
||||
};
|
||||
use lancedb::query::ExecutableQuery;
|
||||
use lancedb::query::Query as LanceDbQuery;
|
||||
use lancedb::query::QueryBase;
|
||||
@@ -18,7 +20,7 @@ use crate::error::NapiErrorExt;
|
||||
use crate::iterator::RecordBatchIterator;
|
||||
use crate::rerankers::Reranker;
|
||||
use crate::rerankers::RerankerCallbacks;
|
||||
use crate::util::{parse_distance_type, parse_fts_query};
|
||||
use crate::util::parse_distance_type;
|
||||
|
||||
#[napi]
|
||||
pub struct Query {
|
||||
@@ -38,51 +40,8 @@ impl Query {
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub fn full_text_search(&mut self, query: napi::JsUnknown) -> napi::Result<()> {
|
||||
let query = unsafe { query.cast::<napi::JsObject>() };
|
||||
let query = if let Some(query_text) = query.get::<_, String>("query").transpose() {
|
||||
let mut query_text = query_text?;
|
||||
let columns = query.get::<_, Option<Vec<String>>>("columns")?.flatten();
|
||||
|
||||
let is_phrase =
|
||||
query_text.len() >= 2 && query_text.starts_with('"') && query_text.ends_with('"');
|
||||
let is_multi_match = columns.as_ref().map(|cols| cols.len() > 1).unwrap_or(false);
|
||||
|
||||
if is_phrase {
|
||||
// Remove the surrounding quotes for phrase queries
|
||||
query_text = query_text[1..query_text.len() - 1].to_string();
|
||||
}
|
||||
|
||||
let query: FtsQuery = match (is_phrase, is_multi_match) {
|
||||
(false, _) => MatchQuery::new(query_text).into(),
|
||||
(true, false) => PhraseQuery::new(query_text).into(),
|
||||
(true, true) => {
|
||||
return Err(napi::Error::from_reason(
|
||||
"Phrase queries cannot be used with multiple columns.",
|
||||
));
|
||||
}
|
||||
};
|
||||
let mut query = FullTextSearchQuery::new_query(query);
|
||||
if let Some(cols) = columns {
|
||||
if !cols.is_empty() {
|
||||
query = query.with_columns(&cols).map_err(|e| {
|
||||
napi::Error::from_reason(format!(
|
||||
"Failed to set full text search columns: {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
}
|
||||
}
|
||||
query
|
||||
} else if let Some(query) = query.get::<_, napi::JsObject>("query")? {
|
||||
let query = parse_fts_query(&query)?;
|
||||
FullTextSearchQuery::new_query(query)
|
||||
} else {
|
||||
return Err(napi::Error::from_reason(
|
||||
"Invalid full text search query object".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
pub fn full_text_search(&mut self, query: napi::JsObject) -> napi::Result<()> {
|
||||
let query = parse_fts_query(query)?;
|
||||
self.inner = self.inner.clone().full_text_search(query);
|
||||
Ok(())
|
||||
}
|
||||
@@ -243,51 +202,8 @@ impl VectorQuery {
|
||||
}
|
||||
|
||||
#[napi]
|
||||
pub fn full_text_search(&mut self, query: napi::JsUnknown) -> napi::Result<()> {
|
||||
let query = unsafe { query.cast::<napi::JsObject>() };
|
||||
let query = if let Some(query_text) = query.get::<_, String>("query").transpose() {
|
||||
let mut query_text = query_text?;
|
||||
let columns = query.get::<_, Option<Vec<String>>>("columns")?.flatten();
|
||||
|
||||
let is_phrase =
|
||||
query_text.len() >= 2 && query_text.starts_with('"') && query_text.ends_with('"');
|
||||
let is_multi_match = columns.as_ref().map(|cols| cols.len() > 1).unwrap_or(false);
|
||||
|
||||
if is_phrase {
|
||||
// Remove the surrounding quotes for phrase queries
|
||||
query_text = query_text[1..query_text.len() - 1].to_string();
|
||||
}
|
||||
|
||||
let query: FtsQuery = match (is_phrase, is_multi_match) {
|
||||
(false, _) => MatchQuery::new(query_text).into(),
|
||||
(true, false) => PhraseQuery::new(query_text).into(),
|
||||
(true, true) => {
|
||||
return Err(napi::Error::from_reason(
|
||||
"Phrase queries cannot be used with multiple columns.",
|
||||
));
|
||||
}
|
||||
};
|
||||
let mut query = FullTextSearchQuery::new_query(query);
|
||||
if let Some(cols) = columns {
|
||||
if !cols.is_empty() {
|
||||
query = query.with_columns(&cols).map_err(|e| {
|
||||
napi::Error::from_reason(format!(
|
||||
"Failed to set full text search columns: {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
}
|
||||
}
|
||||
query
|
||||
} else if let Some(query) = query.get::<_, napi::JsObject>("query")? {
|
||||
let query = parse_fts_query(&query)?;
|
||||
FullTextSearchQuery::new_query(query)
|
||||
} else {
|
||||
return Err(napi::Error::from_reason(
|
||||
"Invalid full text search query object".to_string(),
|
||||
));
|
||||
};
|
||||
|
||||
pub fn full_text_search(&mut self, query: napi::JsObject) -> napi::Result<()> {
|
||||
let query = parse_fts_query(query)?;
|
||||
self.inner = self.inner.clone().full_text_search(query);
|
||||
Ok(())
|
||||
}
|
||||
@@ -376,3 +292,118 @@ impl VectorQuery {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[napi]
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct JsFullTextQuery {
|
||||
pub(crate) inner: FtsQuery,
|
||||
}
|
||||
|
||||
#[napi]
|
||||
impl JsFullTextQuery {
|
||||
#[napi(factory)]
|
||||
pub fn match_query(
|
||||
query: String,
|
||||
column: String,
|
||||
boost: f64,
|
||||
fuzziness: Option<u32>,
|
||||
max_expansions: u32,
|
||||
) -> napi::Result<Self> {
|
||||
Ok(Self {
|
||||
inner: MatchQuery::new(query)
|
||||
.with_column(Some(column))
|
||||
.with_boost(boost as f32)
|
||||
.with_fuzziness(fuzziness)
|
||||
.with_max_expansions(max_expansions as usize)
|
||||
.into(),
|
||||
})
|
||||
}
|
||||
|
||||
#[napi(factory)]
|
||||
pub fn phrase_query(query: String, column: String) -> napi::Result<Self> {
|
||||
Ok(Self {
|
||||
inner: PhraseQuery::new(query).with_column(Some(column)).into(),
|
||||
})
|
||||
}
|
||||
|
||||
#[napi(factory)]
|
||||
pub fn boost_query(
|
||||
positive: &JsFullTextQuery,
|
||||
negative: &JsFullTextQuery,
|
||||
negative_boost: Option<f64>,
|
||||
) -> napi::Result<Self> {
|
||||
Ok(Self {
|
||||
inner: BoostQuery::new(
|
||||
positive.inner.clone(),
|
||||
negative.inner.clone(),
|
||||
negative_boost.map(|v| v as f32),
|
||||
)
|
||||
.into(),
|
||||
})
|
||||
}
|
||||
|
||||
#[napi(factory)]
|
||||
pub fn multi_match_query(
|
||||
query: String,
|
||||
columns: Vec<String>,
|
||||
boosts: Option<Vec<f64>>,
|
||||
) -> napi::Result<Self> {
|
||||
let q = match boosts {
|
||||
Some(boosts) => MultiMatchQuery::try_new_with_boosts(
|
||||
query,
|
||||
columns,
|
||||
boosts.into_iter().map(|v| v as f32).collect(),
|
||||
),
|
||||
None => MultiMatchQuery::try_new(query, columns),
|
||||
}
|
||||
.map_err(|e| {
|
||||
napi::Error::from_reason(format!("Failed to create multi match query: {}", e))
|
||||
})?;
|
||||
|
||||
Ok(Self { inner: q.into() })
|
||||
}
|
||||
}
|
||||
|
||||
fn parse_fts_query(query: napi::JsObject) -> napi::Result<FullTextSearchQuery> {
|
||||
if let Ok(Some(query)) = query.get::<_, &JsFullTextQuery>("query") {
|
||||
Ok(FullTextSearchQuery::new_query(query.inner.clone()))
|
||||
} else if let Ok(Some(query_text)) = query.get::<_, String>("query") {
|
||||
let mut query_text = query_text;
|
||||
let columns = query.get::<_, Option<Vec<String>>>("columns")?.flatten();
|
||||
|
||||
let is_phrase =
|
||||
query_text.len() >= 2 && query_text.starts_with('"') && query_text.ends_with('"');
|
||||
let is_multi_match = columns.as_ref().map(|cols| cols.len() > 1).unwrap_or(false);
|
||||
|
||||
if is_phrase {
|
||||
// Remove the surrounding quotes for phrase queries
|
||||
query_text = query_text[1..query_text.len() - 1].to_string();
|
||||
}
|
||||
|
||||
let query: FtsQuery = match (is_phrase, is_multi_match) {
|
||||
(false, _) => MatchQuery::new(query_text).into(),
|
||||
(true, false) => PhraseQuery::new(query_text).into(),
|
||||
(true, true) => {
|
||||
return Err(napi::Error::from_reason(
|
||||
"Phrase queries cannot be used with multiple columns.",
|
||||
));
|
||||
}
|
||||
};
|
||||
let mut query = FullTextSearchQuery::new_query(query);
|
||||
if let Some(cols) = columns {
|
||||
if !cols.is_empty() {
|
||||
query = query.with_columns(&cols).map_err(|e| {
|
||||
napi::Error::from_reason(format!(
|
||||
"Failed to set full text search columns: {}",
|
||||
e
|
||||
))
|
||||
})?;
|
||||
}
|
||||
}
|
||||
Ok(query)
|
||||
} else {
|
||||
Err(napi::Error::from_reason(
|
||||
"Invalid full text search query object".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
// SPDX-License-Identifier: Apache-2.0
|
||||
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||||
|
||||
use lancedb::index::scalar::{BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, PhraseQuery};
|
||||
use lancedb::DistanceType;
|
||||
|
||||
pub fn parse_distance_type(distance_type: impl AsRef<str>) -> napi::Result<DistanceType> {
|
||||
@@ -16,144 +15,3 @@ pub fn parse_distance_type(distance_type: impl AsRef<str>) -> napi::Result<Dista
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_fts_query(query: &napi::JsObject) -> napi::Result<FtsQuery> {
|
||||
let query_type = query
|
||||
.get_property_names()?
|
||||
.get_element::<napi::JsString>(0)?;
|
||||
let query_type = query_type.into_utf8()?.into_owned()?;
|
||||
let query_value =
|
||||
query
|
||||
.get::<_, napi::JsObject>(&query_type)?
|
||||
.ok_or(napi::Error::from_reason(format!(
|
||||
"query value {} not found",
|
||||
query_type
|
||||
)))?;
|
||||
|
||||
match query_type.as_str() {
|
||||
"match" => {
|
||||
let column = query_value
|
||||
.get_property_names()?
|
||||
.get_element::<napi::JsString>(0)?
|
||||
.into_utf8()?
|
||||
.into_owned()?;
|
||||
let params =
|
||||
query_value
|
||||
.get::<_, napi::JsObject>(&column)?
|
||||
.ok_or(napi::Error::from_reason(format!(
|
||||
"column {} not found",
|
||||
column
|
||||
)))?;
|
||||
|
||||
let query = params
|
||||
.get::<_, napi::JsString>("query")?
|
||||
.ok_or(napi::Error::from_reason("query not found"))?
|
||||
.into_utf8()?
|
||||
.into_owned()?;
|
||||
let boost = params
|
||||
.get::<_, napi::JsNumber>("boost")?
|
||||
.ok_or(napi::Error::from_reason("boost not found"))?
|
||||
.get_double()? as f32;
|
||||
let fuzziness = params
|
||||
.get::<_, napi::JsNumber>("fuzziness")?
|
||||
.map(|f| f.get_uint32())
|
||||
.transpose()?;
|
||||
let max_expansions = params
|
||||
.get::<_, napi::JsNumber>("max_expansions")?
|
||||
.ok_or(napi::Error::from_reason("max_expansions not found"))?
|
||||
.get_uint32()? as usize;
|
||||
|
||||
let query = MatchQuery::new(query)
|
||||
.with_column(Some(column))
|
||||
.with_boost(boost)
|
||||
.with_fuzziness(fuzziness)
|
||||
.with_max_expansions(max_expansions);
|
||||
Ok(query.into())
|
||||
}
|
||||
|
||||
"match_phrase" => {
|
||||
let column = query_value
|
||||
.get_property_names()?
|
||||
.get_element::<napi::JsString>(0)?
|
||||
.into_utf8()?
|
||||
.into_owned()?;
|
||||
let query = query_value
|
||||
.get::<_, napi::JsString>(&column)?
|
||||
.ok_or(napi::Error::from_reason(format!(
|
||||
"column {} not found",
|
||||
column
|
||||
)))?
|
||||
.into_utf8()?
|
||||
.into_owned()?;
|
||||
|
||||
let query = PhraseQuery::new(query).with_column(Some(column));
|
||||
Ok(query.into())
|
||||
}
|
||||
|
||||
"boost" => {
|
||||
let positive = query_value
|
||||
.get::<_, napi::JsObject>("positive")?
|
||||
.ok_or(napi::Error::from_reason("positive not found"))?;
|
||||
|
||||
let negative = query_value
|
||||
.get::<_, napi::JsObject>("negative")?
|
||||
.ok_or(napi::Error::from_reason("negative not found"))?;
|
||||
let negative_boost = query_value
|
||||
.get::<_, napi::JsNumber>("negative_boost")?
|
||||
.ok_or(napi::Error::from_reason("negative_boost not found"))?
|
||||
.get_double()? as f32;
|
||||
|
||||
let positive = parse_fts_query(&positive)?;
|
||||
let negative = parse_fts_query(&negative)?;
|
||||
let query = BoostQuery::new(positive, negative, Some(negative_boost));
|
||||
Ok(query.into())
|
||||
}
|
||||
|
||||
"multi_match" => {
|
||||
let query = query_value
|
||||
.get::<_, napi::JsString>("query")?
|
||||
.ok_or(napi::Error::from_reason("query not found"))?
|
||||
.into_utf8()?
|
||||
.into_owned()?;
|
||||
let columns_array = query_value
|
||||
.get::<_, napi::JsTypedArray>("columns")?
|
||||
.ok_or(napi::Error::from_reason("columns not found"))?;
|
||||
let columns_num = columns_array.get_array_length()?;
|
||||
let mut columns = Vec::with_capacity(columns_num as usize);
|
||||
for i in 0..columns_num {
|
||||
let column = columns_array
|
||||
.get_element::<napi::JsString>(i)?
|
||||
.into_utf8()?
|
||||
.into_owned()?;
|
||||
columns.push(column);
|
||||
}
|
||||
let boost_array = query_value
|
||||
.get::<_, napi::JsTypedArray>("boost")?
|
||||
.ok_or(napi::Error::from_reason("boost not found"))?;
|
||||
if boost_array.get_array_length()? != columns_num {
|
||||
return Err(napi::Error::from_reason(format!(
|
||||
"boost array length ({}) does not match columns length ({})",
|
||||
boost_array.get_array_length()?,
|
||||
columns_num
|
||||
)));
|
||||
}
|
||||
let mut boost = Vec::with_capacity(columns_num as usize);
|
||||
for i in 0..columns_num {
|
||||
let b = boost_array.get_element::<napi::JsNumber>(i)?.get_double()? as f32;
|
||||
boost.push(b);
|
||||
}
|
||||
|
||||
let query =
|
||||
MultiMatchQuery::try_new_with_boosts(query, columns, boost).map_err(|e| {
|
||||
napi::Error::from_reason(format!("Error creating MultiMatchQuery: {}", e))
|
||||
})?;
|
||||
|
||||
Ok(query.into())
|
||||
}
|
||||
|
||||
_ => Err(napi::Error::from_reason(format!(
|
||||
"Unsupported query type: {}",
|
||||
query_type
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user