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

<!-- lance-gatekeeper-fix:v1 agent=97d6925d4af7cd09767da8b552c9bcf3
generation=1 -->

Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com>

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<VectorQuery> {
         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::<Float32Type>()
+                .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<Self> {
-        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<VectorQuery>| {
+            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<dyn Array>; 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();

rust/lancedb/src/remote/table.rs

@@ -1180,11 +1180,14 @@ impl<S: HttpSend> RemoteTable<S> {
                             .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::<Result<Vec<_>>>()?,
                     ))
                 }
                 _ => 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<String> {
+                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(