From a8baecb3065fd75b8be4b47f16f76f60de807934 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Tue, 6 Oct 2026 16:26:25 -0700 Subject: [PATCH] fix: reject non-finite query vectors without panicking (#4410) Remote vector serialization unwrapped `serde_json::Number::from_f64`, which returns `None` for NaN and infinity, causing a Tokio worker panic. Validate vectors after conversion to Float32 in the shared `nearest_to` / `add_query_vector` builder so local and remote queries consistently return `Error::InvalidInput` (mapped to `ValueError` by the Python bindings). This also rejects finite Float64 inputs that overflow Float32. Make remote serialization fallible to protect requests constructed directly, including single-vector queries and both remote batch-query paths. Validation: - All 116 query-related Rust tests passed with the `remote` feature, including the three new regression cases. - The three regression cases fail with the original production code and pass with the fix. Coverage includes NaN, positive and negative infinity, Float16/32/64 Arrow inputs, Float64 overflow, and rejection before any HTTP request. - `cargo fmt --all -- --check`, `cargo check --profile ci --locked -p lancedb --features remote --tests --examples`, and `cargo clippy --profile ci --locked -p lancedb --features remote --tests --examples -- -D warnings` passed. Fixes #4375 Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> --- rust/lancedb/src/query.rs | 93 ++++++++++++++++++++++++++++++-- rust/lancedb/src/remote/table.rs | 52 ++++++++++++++++-- 2 files changed, 136 insertions(+), 9 deletions(-) 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(