From 3d562a7c28ec9fd54401792bf3c3a55e07d7e707 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 05:16:03 +0000 Subject: [PATCH] fix(python): reject NaNs in multivector columns --- python/python/lancedb/table.py | 161 ++++++++++++++---- python/python/tests/test_table.py | 13 ++ python/python/tests/test_util.py | 77 +++++++++ rust/lancedb/src/table/add_data.rs | 58 +++++++ .../src/table/datafusion/reject_nan.rs | 134 +++++++++++---- 5 files changed, 374 insertions(+), 69 deletions(-) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index ae36bac7a..e46187a49 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -4035,7 +4035,10 @@ def _handle_bad_vectors( for vector_column in vector_columns: dim = vector_column["expected_dim"] if target_schema is not None and dim is None: - dim = _infer_vector_dim(batch[vector_column["name"]]) + dim = _infer_vector_column_dim( + batch[vector_column["name"]], + vector_column["is_multivector"], + ) pending_dims.append(vector_column) batch = _handle_bad_vector_column( batch, @@ -4044,11 +4047,13 @@ def _handle_bad_vectors( fill_value=fill_value, expected_dim=dim, expected_value_type=vector_column["expected_value_type"], + is_multivector=vector_column["is_multivector"], ) for vector_column in pending_dims: if vector_column["expected_dim"] is None: - vector_column["expected_dim"] = _infer_vector_dim( - batch[vector_column["name"]] + vector_column["expected_dim"] = _infer_vector_column_dim( + batch[vector_column["name"]], + vector_column["is_multivector"], ) if batch.schema.equals(output_schema, check_metadata=True): yield batch @@ -4074,22 +4079,30 @@ def _find_vector_columns( if target_schema is None: vector_columns = [] for field in reader_schema: - named_vector_col = ( - _is_list_like(field.type) - and pa.types.is_floating(field.type.value_type) - and field.name == VECTOR_COLUMN_NAME + is_multivector = _is_multivector_type(field.type) + is_fixed_multivector = is_multivector and pa.types.is_fixed_size_list( + field.type.value_type ) + named_vector_col = ( + _is_float_vector_type(field.type) or is_multivector + ) and field.name == VECTOR_COLUMN_NAME likely_vector_col = ( pa.types.is_fixed_size_list(field.type) and pa.types.is_floating(field.type.value_type) and (field.type.list_size >= 10) ) - if named_vector_col or likely_vector_col: + if named_vector_col or likely_vector_col or is_fixed_multivector: + vector_type = field.type.value_type if is_multivector else field.type vector_columns.append( { "name": field.name, - "expected_dim": None, - "expected_value_type": None, + "expected_dim": ( + vector_type.list_size + if pa.types.is_fixed_size_list(vector_type) + else None + ), + "expected_value_type": vector_type.value_type, + "is_multivector": is_multivector, } ) return vector_columns @@ -4103,9 +4116,8 @@ def _find_vector_columns( for field in target_schema: if field.name not in reader_column_names: continue - if not _is_list_like(field.type) or not pa.types.is_floating( - field.type.value_type - ): + is_multivector = _is_multivector_type(field.type) + if not _is_float_vector_type(field.type) and not is_multivector: continue reader_field = reader_schema.field(field.name) @@ -4120,16 +4132,18 @@ def _find_vector_columns( and reader_field.type.list_size >= 10 ) - if named_vector_col or typed_fixed_vector_col: + if named_vector_col or typed_fixed_vector_col or is_multivector: + vector_type = field.type.value_type if is_multivector else field.type vector_columns.append( { "name": field.name, "expected_dim": ( - field.type.list_size - if pa.types.is_fixed_size_list(field.type) + vector_type.list_size + if pa.types.is_fixed_size_list(vector_type) else None ), - "expected_value_type": field.type.value_type, + "expected_value_type": vector_type.value_type, + "is_multivector": is_multivector, } ) @@ -4180,6 +4194,7 @@ def _handle_bad_vector_column( fill_value: float = 0.0, expected_dim: Optional[int] = None, expected_value_type: Optional[pa.DataType] = None, + is_multivector: bool = False, ) -> pa.RecordBatch: """ Ensure that the vector column exists and has type fixed_size_list(float) @@ -4200,6 +4215,8 @@ def _handle_bad_vector_column( vec_arr = data[vector_column_name] if not _is_list_like(vec_arr.type): return data + if is_multivector and not _is_multivector_type(vec_arr.type): + return data if ( expected_dim is not None @@ -4216,12 +4233,21 @@ def _handle_bad_vector_column( vec_arr = pa.array(vec_arr.to_pylist(), type=pa.list_(expected_value_type)) data = data.set_column(position, vector_column_name, vec_arr) - if pa.types.is_floating(vec_arr.type.value_type): + if is_multivector or pa.types.is_floating(vec_arr.type.value_type): has_nan = has_nan_values(vec_arr) else: has_nan = pa.array([False] * len(vec_arr)) - if expected_dim is not None: + if is_multivector: + dim = ( + expected_dim + if expected_dim is not None + else _infer_vector_column_dim(vec_arr, True) + ) + if dim is None: + return data + has_wrong_dim = _multivector_has_wrong_dim(vec_arr, dim) + elif expected_dim is not None: dim = expected_dim elif pa.types.is_fixed_size_list(vec_arr.type): dim = vec_arr.type.list_size @@ -4230,15 +4256,16 @@ def _handle_bad_vector_column( if dim is None: return data - is_null = pc.is_null(vec_arr) - # pc.list_value_length returns null for null list entries, so - # pc.not_equal(null, dim) also returns null. Use or_kleene so that - # True OR null = True (Kleene three-valued logic), ensuring null vectors - # are counted as wrong-dim. - has_wrong_dim = pc.or_kleene( - is_null, - pc.not_equal(pc.list_value_length(vec_arr), dim), - ) + if not is_multivector: + is_null = pc.is_null(vec_arr) + # pc.list_value_length returns null for null list entries, so + # pc.not_equal(null, dim) also returns null. Use or_kleene so that + # True OR null = True (Kleene three-valued logic), ensuring null vectors + # are counted as wrong-dim. + has_wrong_dim = pc.or_kleene( + is_null, + pc.not_equal(pc.list_value_length(vec_arr), dim), + ) has_bad_vectors = pc.any(has_nan).as_py() or pc.any(has_wrong_dim).as_py() @@ -4273,7 +4300,10 @@ def _handle_bad_vector_column( raise ValueError( "`fill_value` must not be None if `on_bad_vectors` is 'fill'" ) - vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value) + if is_multivector: + vec_arr = _fill_bad_multivector_values(vec_arr, dim, fill_value) + else: + vec_arr = _fill_bad_vector_values(vec_arr, dim, fill_value) else: raise ValueError(f"Invalid value for on_bad_vectors: {on_bad_vectors}") @@ -4325,17 +4355,60 @@ def _fill_bad_vector_values( return filled.cast(arr.type) +def _fill_bad_multivector_values( + arr: Union[pa.Array, pa.ChunkedArray], dim: int, fill_value: float +) -> pa.Array: + if not isinstance(arr, pa.ChunkedArray): + arr = pa.chunked_array([arr]) + arr = arr.combine_chunks() + + filled_vectors = _fill_bad_vector_values(arr.values, dim, fill_value) + parent_nulls = pc.is_null(arr) + if pa.types.is_large_list(arr.type): + filled = pa.LargeListArray.from_arrays( + arr.offsets, filled_vectors, mask=parent_nulls + ) + else: + filled = pa.ListArray.from_arrays( + arr.offsets, filled_vectors, mask=parent_nulls + ) + return filled.cast(arr.type) + + +def _multivector_has_wrong_dim( + arr: Union[pa.Array, pa.ChunkedArray], dim: int +) -> pa.BooleanArray: + if isinstance(arr, pa.ChunkedArray): + results = [_multivector_has_wrong_dim(chunk, dim) for chunk in arr.chunks] + return pa.concat_arrays(results) if results else pa.array([], type=pa.bool_()) + + vectors = arr.flatten() + vector_is_wrong = pc.or_kleene( + pc.is_null(vectors), + pc.not_equal(pc.list_value_length(vectors), dim), + ) + parent_indices = pc.list_parent_indices(arr) + wrong_parent_indices = pc.unique(pc.filter(parent_indices, vector_is_wrong)) + indices = pa.array(range(len(arr)), type=pa.uint32()) + return pc.or_(pc.is_null(arr), pc.is_in(indices, wrong_parent_indices)) + + def has_nan_values(arr: Union[pa.ListArray, pa.ChunkedArray]) -> pa.BooleanArray: if isinstance(arr, pa.ChunkedArray): - values = pa.chunked_array([chunk.flatten() for chunk in arr.chunks]) - else: - values = arr.flatten() - if pa.types.is_float16(values.type): + results = [has_nan_values(chunk) for chunk in arr.chunks] + return pa.concat_arrays(results) if results else pa.array([], type=pa.bool_()) + + values = arr.flatten() + if _is_list_like(values.type): + values_has_nan = has_nan_values(values) + elif pa.types.is_float16(values.type): # is_nan isn't yet implemented for f16, so we cast to f32 # https://github.com/apache/arrow/issues/45083 values_has_nan = pc.is_nan(values.cast(pa.float32())) - else: + elif pa.types.is_floating(values.type): values_has_nan = pc.is_nan(values) + else: + return pa.array([False] * len(arr)) values_indices = pc.list_parent_indices(arr) has_nan_indices = pc.unique(pc.filter(values_indices, values_has_nan)) indices = pa.array(range(len(arr)), type=pa.uint32()) @@ -4350,6 +4423,16 @@ def _is_list_like(data_type: pa.DataType) -> bool: ) +def _is_float_vector_type(data_type: pa.DataType) -> bool: + return _is_list_like(data_type) and pa.types.is_floating(data_type.value_type) + + +def _is_multivector_type(data_type: pa.DataType) -> bool: + return ( + pa.types.is_list(data_type) or pa.types.is_large_list(data_type) + ) and _is_float_vector_type(data_type.value_type) + + def _merge_metadata(*metadata_dicts: Optional[dict]) -> dict: merged = {} for metadata in metadata_dicts: @@ -4441,6 +4524,16 @@ def _infer_vector_dim(arr: Union[pa.Array, pa.ChunkedArray]) -> Optional[int]: return pc.mode(lengths)[0].as_py()["mode"] +def _infer_vector_column_dim( + arr: Union[pa.Array, pa.ChunkedArray], is_multivector: bool +) -> Optional[int]: + if not is_multivector: + return _infer_vector_dim(arr) + if isinstance(arr, pa.ChunkedArray): + arr = arr.combine_chunks() + return _infer_vector_dim(arr.flatten()) + + def _validate_schema(schema: pa.Schema): """ Make sure the metadata is valid utf8 diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 069527b21..457fedbab 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -1710,6 +1710,19 @@ def test_create_with_nans(mem_db: DBConnection): assert np.allclose(filled_vectors[22.0], np.array([5.0, 0.0])) +def test_create_with_nans_in_multivectors(mem_db: DBConnection): + multivector_type = pa.list_(pa.list_(pa.float32(), 128)) + schema = pa.schema( + [pa.field("filename", pa.string()), pa.field("vector", multivector_type)] + ) + vector = [0.1] * 128 + vector[-1] = np.nan + data = [{"filename": "img1.jpg", "vector": [vector]}] + + with pytest.raises(RuntimeError, match="Vector column 'vector' has NaNs"): + mem_db.create_table("nan_multivector", data=data, schema=schema) + + def test_add_with_nans(mem_db: DBConnection): schema = pa.schema( [ diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index a9b66b2dd..824d5fd3f 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -400,6 +400,83 @@ def test_handle_bad_vectors_nan(on_bad_vectors): assert output["vector"].combine_chunks() == expected +@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) +def test_handle_bad_multivectors_nan(on_bad_vectors): + multivector_type = pa.list_(pa.list_(pa.float32(), 2)) + vectors = pa.array( + [ + [[1.0, float("nan")], [2.0, 3.0]], + [[4.0, 5.0]], + ], + type=multivector_type, + ) + data = pa.table({"vector": vectors}) + + if on_bad_vectors == "error": + with pytest.raises(ValueError, match="Vector column 'vector' has NaNs"): + _handle_bad_vectors(data.to_reader()).read_all() + return + + output = _handle_bad_vectors( + data.to_reader(), + on_bad_vectors=on_bad_vectors, + fill_value=42.0, + ).read_all() + + if on_bad_vectors == "drop": + expected = pa.array([[[4.0, 5.0]]], type=multivector_type) + elif on_bad_vectors == "fill": + expected = pa.array( + [[[1.0, 42.0], [2.0, 3.0]], [[4.0, 5.0]]], + type=multivector_type, + ) + else: + expected = pa.array([None, [[4.0, 5.0]]], type=multivector_type) + + assert output["vector"].combine_chunks() == expected + + +@pytest.mark.parametrize("on_bad_vectors", ["error", "drop", "fill", "null"]) +def test_handle_bad_variable_multivectors(on_bad_vectors): + target_type = pa.list_(pa.list_(pa.float32(), 2)) + vectors = pa.array( + [ + [[1.0, float("nan")], [2.0, 3.0]], + [[4.0]], + [[5.0, 6.0]], + ] + ) + data = pa.table({"vector": vectors}) + + if on_bad_vectors == "error": + with pytest.raises(ValueError, match="variable length vectors"): + _handle_bad_vectors( + data.to_reader(), + target_schema=pa.schema({"vector": target_type}), + ).read_all() + return + + output = _handle_bad_vectors( + data.to_reader(), + on_bad_vectors=on_bad_vectors, + fill_value=42.0, + target_schema=pa.schema({"vector": target_type}), + ).read_all() + + if on_bad_vectors == "drop": + expected = [[[5.0, 6.0]]] + elif on_bad_vectors == "fill": + expected = [ + [[1.0, 42.0], [2.0, 3.0]], + [[4.0, 42.0]], + [[5.0, 6.0]], + ] + else: + expected = [None, None, [[5.0, 6.0]]] + + assert output["vector"].combine_chunks().to_pylist() == expected + + def test_handle_bad_vectors_noop(): # ChunkedArray should be preserved as-is vector = pa.chunked_array( diff --git a/rust/lancedb/src/table/add_data.rs b/rust/lancedb/src/table/add_data.rs index 11ba43dd6..601d40a51 100644 --- a/rust/lancedb/src/table/add_data.rs +++ b/rust/lancedb/src/table/add_data.rs @@ -258,6 +258,7 @@ mod tests { FixedSizeListArray, Float32Array, Int32Array, LargeStringArray, ListArray, RecordBatch, RecordBatchIterator, record_batch, }; + use arrow_buffer::OffsetBuffer; use arrow_schema::{ArrowError, DataType, Field, Schema}; use futures::TryStreamExt; use lance::dataset::{WriteMode, WriteParams}; @@ -885,6 +886,63 @@ mod tests { assert_eq!(row_count, 1); } + #[tokio::test] + async fn test_add_rejects_nan_multivectors() { + let vector_type = + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4); + let schema = Arc::new(Schema::new(vec![Field::new( + "embedding", + DataType::List(Arc::new(Field::new("item", vector_type.clone(), true))), + false, + )])); + + let db = connect("memory://").execute().await.unwrap(); + let table = db + .create_empty_table("nan_multivector_test", schema.clone()) + .execute() + .await + .unwrap(); + + let vectors = FixedSizeListArray::try_new( + Arc::new(Field::new("item", DataType::Float32, true)), + 4, + Arc::new(Float32Array::from(vec![ + 0.1, + 0.2, + 0.3, + 0.4, + 0.5, + f32::NAN, + 0.7, + 0.8, + ])), + None, + ) + .unwrap(); + let multivectors = ListArray::try_new( + Arc::new(Field::new("item", vector_type, true)), + OffsetBuffer::from_lengths([2]), + Arc::new(vectors), + None, + ) + .unwrap(); + let batch = RecordBatch::try_new(schema, vec![Arc::new(multivectors)]).unwrap(); + + let err = table.add(batch.clone()).execute().await.unwrap_err(); + assert!( + err.to_string().contains("NaN"), + "Expected error mentioning NaN values, but got: {err:?}" + ); + + table + .add(batch) + .on_nan_vectors(NaNVectorBehavior::Keep) + .execute() + .await + .unwrap(); + assert_eq!(table.count_rows(None).await.unwrap(), 1); + } + #[tokio::test] async fn test_add_subschema() { let data = record_batch!(("id", Int64, [4, 5]), ("text", Utf8, ["foo", "bar"])).unwrap(); diff --git a/rust/lancedb/src/table/datafusion/reject_nan.rs b/rust/lancedb/src/table/datafusion/reject_nan.rs index 084c9722b..4798532ce 100644 --- a/rust/lancedb/src/table/datafusion/reject_nan.rs +++ b/rust/lancedb/src/table/datafusion/reject_nan.rs @@ -5,7 +5,7 @@ use std::sync::{Arc, LazyLock}; -use arrow_array::{Array, FixedSizeListArray}; +use arrow_array::{Array, FixedSizeListArray, ListArray}; use arrow_schema::{DataType, Field, FieldRef}; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; @@ -19,16 +19,20 @@ use crate::{Error, Result}; static REJECT_NAN_UDF: LazyLock> = LazyLock::new(|| Arc::new(datafusion_expr::ScalarUDF::from(RejectNanUdf::new()))); -/// Returns true if the field is a vector column: FixedSizeList. +/// Returns true if the field is a vector or multivector column. fn is_vector_field(field: &Field) -> bool { - if let DataType::FixedSizeList(child, _) = field.data_type() { - matches!( - child.data_type(), - DataType::Float16 | DataType::Float32 | DataType::Float64 - ) - } else { - false + fn is_vector_data_type(data_type: &DataType) -> bool { + match data_type { + DataType::FixedSizeList(child, _) => matches!( + child.data_type(), + DataType::Float16 | DataType::Float32 | DataType::Float64 + ), + DataType::List(child) => is_vector_data_type(child.data_type()), + _ => false, + } } + + is_vector_data_type(field.data_type()) } /// Wraps the input plan with a projection that checks vector columns for NaN values. @@ -69,8 +73,8 @@ pub fn reject_nan_vectors(input: Arc) -> Result datafusion_common::Result<()> { - let fsl = array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - datafusion_common::DataFusionError::Internal( - "reject_nan expected FixedSizeList".to_string(), - ) - })?; - - // Only inspect elements that are both in a valid parent row and non-null - // themselves. Values backing null parent rows or null child elements may - // contain garbage (including NaN) per the Arrow spec. - let has_nan = (0..fsl.len()).filter(|i| fsl.is_valid(*i)).any(|i| { - let row = fsl.value(i); - match row.data_type() { - DataType::Float16 => row + fn contains_nan(array: &dyn Array) -> datafusion_common::Result { + match array.data_type() { + DataType::Float16 => Ok(array .as_any() .downcast_ref::() .unwrap() .iter() - .any(|v| v.is_some_and(|v| v.is_nan())), - DataType::Float32 => row + .any(|v| v.is_some_and(|v| v.is_nan()))), + DataType::Float32 => Ok(array .as_any() .downcast_ref::() .unwrap() .iter() - .any(|v| v.is_some_and(|v| v.is_nan())), - DataType::Float64 => row + .any(|v| v.is_some_and(|v| v.is_nan()))), + DataType::Float64 => Ok(array .as_any() .downcast_ref::() .unwrap() .iter() - .any(|v| v.is_some_and(|v| v.is_nan())), - _ => false, + .any(|v| v.is_some_and(|v| v.is_nan()))), + DataType::FixedSizeList(_, _) => { + let lists = array.as_any().downcast_ref::().unwrap(); + for i in (0..lists.len()).filter(|i| lists.is_valid(*i)) { + if contains_nan(lists.value(i).as_ref())? { + return Ok(true); + } + } + Ok(false) + } + DataType::List(_) => { + let lists = array.as_any().downcast_ref::().unwrap(); + for i in (0..lists.len()).filter(|i| lists.is_valid(*i)) { + if contains_nan(lists.value(i).as_ref())? { + return Ok(true); + } + } + Ok(false) + } + data_type => Err(datafusion_common::DataFusionError::Internal(format!( + "reject_nan expected a vector or multivector, got {data_type}" + ))), } - }); + } - if has_nan { + // Only inspect elements that are both in a valid parent row and non-null + // themselves. Values backing null parent rows or null child elements may + // contain garbage (including NaN) per the Arrow spec. + if contains_nan(array)? { return Err(datafusion_common::DataFusionError::ArrowError( Box::new(arrow_schema::ArrowError::ComputeError( "Vector column contains NaN values".to_string(), @@ -166,6 +180,7 @@ fn check_no_nans(array: &dyn Array) -> datafusion_common::Result<()> { mod tests { use super::*; use arrow_array::Float32Array; + use arrow_buffer::OffsetBuffer; #[test] fn test_passes_clean_vectors() { @@ -191,6 +206,46 @@ mod tests { assert!(check_no_nans(&fsl).is_err()); } + #[test] + fn test_rejects_nan_multivectors() { + let vectors = FixedSizeListArray::try_new( + Arc::new(Field::new("item", DataType::Float32, true)), + 2, + Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, f32::NAN])), + None, + ) + .unwrap(); + let multivectors = ListArray::try_new( + Arc::new(Field::new("item", vectors.data_type().clone(), true)), + OffsetBuffer::from_lengths([2]), + Arc::new(vectors), + None, + ) + .unwrap(); + + assert!(check_no_nans(&multivectors).is_err()); + } + + #[test] + fn test_skips_null_multivector_rows() { + let vectors = FixedSizeListArray::try_new( + Arc::new(Field::new("item", DataType::Float32, true)), + 2, + Arc::new(Float32Array::from(vec![f32::NAN, f32::NAN, 1.0, 2.0])), + None, + ) + .unwrap(); + let multivectors = ListArray::try_new( + Arc::new(Field::new("item", vectors.data_type().clone(), true)), + OffsetBuffer::from_lengths([1, 1]), + Arc::new(vectors), + Some(vec![false, true].into()), + ) + .unwrap(); + + assert!(check_no_nans(&multivectors).is_ok()); + } + #[test] fn test_skips_null_rows() { // Values backing null rows may contain NaN per the Arrow spec. @@ -254,6 +309,15 @@ mod tests { DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float64, true)), 4), false, ))); + assert!(is_vector_field(&Field::new( + "v", + DataType::List(Arc::new(Field::new( + "item", + DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 4,), + true, + ))), + false, + ))); assert!(!is_vector_field(&Field::new("id", DataType::Int32, false))); assert!(!is_vector_field(&Field::new( "v",