diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index fffb703fb..6b20fb123 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -11,7 +11,7 @@ use arrow_array::{ Array, Float16Array, Float32Array, Float64Array, RecordBatch, UInt64Array, cast::AsArray, make_array, - types::{Int64Type, UInt64Type}, + types::{Float32Type, Int64Type, UInt64Type}, }; use arrow_schema::{DataType, SchemaRef}; use datafusion_common::{DataFusionError, Result as DataFusionResult}; @@ -1024,6 +1024,8 @@ impl Query { /// /// By default, there is no embedding model, and the input should be /// vector/slice of floats. + /// NaN and infinite components are rejected with [`Error::InvalidInput`] + /// after conversion to Float32, for both local and remote tables. /// /// If there is only one vector column (a column whose data type is a /// fixed size list of floats) then the column does not need to be specified. @@ -1050,8 +1052,7 @@ impl Query { /// * `vector` - The vector that will be used for search. pub fn nearest_to(self, vector: impl IntoQueryVector) -> Result { let mut vector_query = self.into_vector(); - let query_vector = vector.to_query_vector(&DataType::Float32, "default")?; - vector_query.request.query_vector.push(query_vector); + vector_query.request.add_query_vector(vector)?; if vector_query.request.base.limit.is_none() { vector_query.request.base.limit = Some(DEFAULT_TOP_K); @@ -1154,6 +1155,23 @@ impl Default for VectorQueryRequest { } impl VectorQueryRequest { + fn add_query_vector(&mut self, vector: impl IntoQueryVector) -> Result<()> { + let vector = vector.to_query_vector(&DataType::Float32, "default")?; + let array = + vector + .as_primitive_opt::() + .ok_or_else(|| Error::InvalidInput { + message: "VectorQuery vector must be of type Float32".into(), + })?; + if array.values().iter().any(|value| !value.is_finite()) { + return Err(Error::InvalidInput { + message: "query vector must contain only finite values".into(), + }); + } + self.query_vector.push(vector); + Ok(()) + } + pub fn from_plain_query(query: QueryRequest) -> Self { Self { base: query, @@ -1222,9 +1240,11 @@ impl VectorQuery { /// The output data will contain an additional column `query_index` which /// will contain the index of the query vector that was used to generate the /// result. + /// + /// NaN and infinite components are rejected with [`Error::InvalidInput`] + /// after conversion to Float32, for both local and remote tables. pub fn add_query_vector(mut self, vector: impl IntoQueryVector) -> Result { - let query_vector = vector.to_query_vector(&DataType::Float32, "default")?; - self.request.query_vector.push(query_vector); + self.request.add_query_vector(vector)?; Ok(self) } @@ -2336,6 +2356,69 @@ mod tests { use crate::{Table, connect, database::CreateTableMode, index::Index}; + #[tokio::test] + async fn test_query_vector_non_finite() { + let conn = connect("memory://non_finite_vectors") + .execute() + .await + .unwrap(); + let table = conn + .create_table("my_table", make_test_batches()) + .execute() + .await + .unwrap(); + let assert_invalid = |result: Result| { + let err = result.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. })); + assert!(err.to_string().contains("only finite values")); + }; + + // Float64 values can be finite before conversion but overflow Float32. + for value in [ + f64::NAN, + f64::INFINITY, + f64::NEG_INFINITY, + f64::MAX, + f64::MIN, + ] { + assert_invalid(table.query().nearest_to(&[0.1, value as f32])); + assert_invalid(table.query().nearest_to(vec![0.1, value])); + assert_invalid( + table + .query() + .nearest_to(&[0.1, 0.2]) + .unwrap() + .add_query_vector(vec![0.1, value]), + ); + + let arrays: [Arc; 3] = [ + Arc::new(Float16Array::from(vec![ + f16::from_f32(0.1), + f16::from_f64(value), + ])), + Arc::new(Float32Array::from(vec![0.1, value as f32])), + Arc::new(Float64Array::from(vec![0.1, value])), + ]; + for array in arrays { + assert_invalid(table.query().nearest_to(array.clone())); + assert_invalid( + table + .query() + .nearest_to(&[0.1, 0.2]) + .unwrap() + .add_query_vector(array), + ); + } + } + + table + .query() + .nearest_to(&[f32::MIN, f32::MAX]) + .unwrap() + .add_query_vector(&[0.0, -0.0]) + .unwrap(); + } + #[tokio::test] async fn test_setters_getters() { let batches = make_test_batches(); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 2fc58e1e5..ca8fe644f 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -1180,11 +1180,14 @@ impl RemoteTable { .values() .iter() .map(|v| { - serde_json::Value::Number( - serde_json::Number::from_f64(*v as f64).unwrap(), - ) + serde_json::Number::from_f64(*v as f64) + .map(serde_json::Value::Number) + .ok_or_else(|| Error::InvalidInput { + message: "query vector must contain only finite values" + .into(), + }) }) - .collect(), + .collect::>>()?, )) } _ => Err(Error::InvalidInput { @@ -5892,6 +5895,47 @@ mod tests { assert_eq!(blobs.value(0), b"alpha"); } + #[rstest] + #[case(DEFAULT_SERVER_VERSION.clone())] + #[case(semver::Version::new(0, 2, 0))] + #[tokio::test] + async fn test_query_vector_non_finite(#[case] version: semver::Version) { + let table = + Table::new_with_handler_version("my_table", version, |_| -> http::Response { + panic!("non-finite vectors must be rejected before sending a request") + }); + + for value in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + // Requests built without the query builder must also return an + // error, including a non-finite vector later in a batch. + for batched in [false, true] { + let mut request = table + .query() + .nearest_to(&[0.1, 0.2]) + .unwrap() + .into_request(); + if !batched { + request.query_vector.clear(); + } + request + .query_vector + .push(Arc::new(arrow_array::Float32Array::from(vec![0.1, value]))); + let result = table + .base_table() + .query( + &AnyQuery::VectorQuery(request), + QueryExecutionOptions::default(), + ) + .await; + let Err(err) = result else { + panic!("non-finite query vector unexpectedly succeeded") + }; + assert!(matches!(err, Error::InvalidInput { .. })); + assert!(err.to_string().contains("only finite values")); + } + } + } + #[tokio::test] async fn test_query_vector_default_values() { let expected_data = RecordBatch::try_new(