diff --git a/python/python/tests/test_vector_index.py b/python/python/tests/test_vector_index.py index d200633ce36..a4b6d960be5 100644 --- a/python/python/tests/test_vector_index.py +++ b/python/python/tests/test_vector_index.py @@ -1296,6 +1296,57 @@ def test_multivec_ann(indexed_multivec_dataset: lance.LanceDataset): ) +def test_multivec_search_paths(tmp_path: Path): + vector_type = pa.list_(pa.list_(pa.float32(), 2)) + query = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32) + uri = tmp_path / "multivec_distance.lance" + + indexed_rows = pa.table( + { + "id": pa.array([0, 1], type=pa.int32()), + "vector": pa.array( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[1.0, 0.0], [1.0, 0.0]], + ], + type=vector_type, + ), + } + ) + dataset = lance.write_dataset(indexed_rows, uri) + dataset = dataset.create_index( + "vector", + index_type="IVF_FLAT", + metric="cosine", + num_partitions=1, + ) + + unindexed_rows = pa.table( + { + "id": pa.array([2, 3], type=pa.int32()), + "vector": pa.array( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[-1.0, 0.0], [0.0, -1.0]], + ], + type=vector_type, + ), + } + ) + dataset = lance.write_dataset(unindexed_rows, uri, mode="append") + + nearest = {"column": "vector", "q": query, "k": 4, "metric": "cosine"} + flat = dataset.to_table(columns=["id"], nearest={**nearest, "use_index": False}) + mixed = dataset.to_table(columns=["id"], nearest=nearest) + + dataset.optimize.optimize_indices() + fully_indexed = dataset.to_table(columns=["id"], nearest=nearest, fast_search=True) + + for result in [flat, mixed, fully_indexed]: + assert result["id"].to_pylist() == [0, 2, 1, 3] + np.testing.assert_allclose(result["_distance"].to_numpy(), [0.0, 0.0, 1.0, 2.0]) + + def test_pre_populated_ivf_centroids(dataset, tmp_path: Path): centroids = np.random.randn(5, 128).astype(np.float32) # IVF5 dataset_with_index = dataset.create_index( diff --git a/rust/lance-linalg/src/distance.rs b/rust/lance-linalg/src/distance.rs index 7b92bab8a52..81a80aaacd4 100644 --- a/rust/lance-linalg/src/distance.rs +++ b/rust/lance-linalg/src/distance.rs @@ -343,6 +343,11 @@ impl TryFrom<&str> for DistanceType { } } +/// Computes the additive late-interaction distance from a multivector query. +/// +/// For each query sub-vector, this finds the minimum distance to any stored +/// sub-vector in the row, then sums those minimum distances. Null or empty +/// stored rows produce `NaN`. pub fn multivec_distance( query: &dyn Array, vectors: &ListArray, @@ -363,7 +368,7 @@ pub fn multivec_distance( // and then downcasts the *stored* values to that same type. The dim, null // and length checks prevent a `chunks_exact` panic and, worse, silently // wrong results: a short query yields no sub-vectors and scores every row - // `1.0`, and a null slot is scored from whatever the values buffer holds. + // `0.0`, and a null slot is scored from whatever the values buffer holds. let query_type = query.data_type(); // Which element types have a kernel here at all. `Int8` is a valid vector // element type elsewhere in the stack (`l2_distance_arrow_batch` and its @@ -424,47 +429,37 @@ pub fn multivec_distance( continue; } - let sim = match distance_type { - DistanceType::Hamming => { - let query = query.as_primitive::().values(); - query - .chunks_exact(dim) - .map(|q| { - multivector - .values() - .as_primitive::() - .values() - .chunks_exact(dim) - .map(|v| hamming::hamming(q, v)) - .min_by(|a, b| a.partial_cmp(b).unwrap()) - .unwrap() - }) - .sum() - } + let distance = match distance_type { + DistanceType::Hamming => multivec_distance_impl::( + query, + multivector, + dim, + hamming::hamming, + ), _ => match query.data_type() { DataType::Float16 => multivec_distance_impl::( query, multivector, dim, - distance_type, + distance_type.func(), ), DataType::Float32 => multivec_distance_impl::( query, multivector, dim, - distance_type, + distance_type.func(), ), DataType::Float64 => multivec_distance_impl::( query, multivector, dim, - distance_type, + distance_type.func(), ), _ => unreachable!("missed to check query type"), }, }; - dists.push(1.0 - sim); + dists.push(distance); } } } @@ -475,11 +470,8 @@ fn multivec_distance_impl( query: &dyn Array, multivector: &FixedSizeListArray, dim: usize, - distance_type: DistanceType, -) -> f32 -where - T::Native: L2 + Cosine + Dot, -{ + distance_func: DistanceFunc, +) -> f32 { let query = query.as_primitive::().values(); query .chunks_exact(dim) @@ -489,8 +481,8 @@ where .as_primitive::() .values() .chunks_exact(dim) - .map(|v| 1.0 - distance_type.func()(q, v)) - .max_by(|a, b| a.total_cmp(b)) + .map(|v| distance_func(q, v)) + .min_by(|a, b| a.total_cmp(b)) .unwrap() }) .sum() @@ -506,7 +498,7 @@ mod tests { use arrow_array::types::{Float16Type, Float32Type, Int8Type}; use arrow_array::{Float32Array, Int8Array, ListArray, PrimitiveArray, UInt8Array}; - use arrow_buffer::OffsetBuffer; + use arrow_buffer::{OffsetBuffer, ScalarBuffer}; use arrow_schema::Field; use half::f16; @@ -529,9 +521,17 @@ mod tests { .expect("write x86 runtime feature report"); } - /// Build a single-row `List>` holding one sub-vector. - fn multivec_of(values: Vec, dim: i32) -> ListArray { - let inner = PrimitiveArray::::from_iter_values(values); + /// Build `List>` rows from flattened sub-vector values. + fn multivecs_of(rows: Vec>, dim: i32) -> ListArray { + let lengths = rows + .iter() + .map(|row| { + assert_eq!(row.len() % dim as usize, 0); + row.len() / dim as usize + }) + .collect::>(); + let values = ScalarBuffer::from(rows.into_iter().flatten().collect::>()); + let inner = PrimitiveArray::::new(values, None); let fsl = FixedSizeListArray::try_new( Arc::new(Field::new("item", T::DATA_TYPE, true)), dim, @@ -539,11 +539,16 @@ mod tests { None, ) .unwrap(); - let offsets = OffsetBuffer::from_lengths([1_usize]); + let offsets = OffsetBuffer::from_lengths(lengths); let field = Arc::new(Field::new("item", fsl.data_type().clone(), true)); ListArray::try_new(field, offsets, Arc::new(fsl), None).unwrap() } + /// Build one `List>` row. + fn multivec_of(values: Vec, dim: i32) -> ListArray { + multivecs_of::(vec![values], dim) + } + /// The `(query dtype, distance type)` pre-check and the dispatch must agree. /// `UInt8` is only valid with Hamming, and the float types only with the /// float metrics; a mismatch must be an error rather than a panic in the @@ -608,7 +613,7 @@ mod tests { /// A query length that is not a positive multiple of `dim` is structurally /// invalid: `chunks_exact` would silently drop the tail, and a query shorter - /// than `dim` would yield no sub-vectors at all and score every row `1.0`. + /// than `dim` would yield no sub-vectors at all and score every row `0.0`. #[test] fn test_multivec_distance_rejects_bad_query_length() { let vectors = multivec_of::(vec![1.0, 2.0], 2); @@ -670,24 +675,62 @@ mod tests { ); } - /// The guards must not reject the combinations that do work: `UInt8` with - /// Hamming is the one non-float path through this function. - /// - /// Note the expected value is `1.0 - hamming`, matching what the function - /// computes. Unlike the float paths — which accumulate `1.0 - distance` and - /// so end up with a distance again — the Hamming path accumulates a raw - /// distance, so `1.0 - sim` inverts its ranking. That inversion is - /// pre-existing and out of scope here; this test pins current behavior - /// rather than endorsing it. + /// Each query sub-vector contributes its minimum Hamming distance to the + /// row total. #[test] - fn test_multivec_distance_accepts_u8_hamming() { - let vectors = multivec_of::(vec![0b0000_1111, 0b0000_0000], 2); - let query: Arc = Arc::new(UInt8Array::from(vec![0b0000_1111_u8, 0b0000_0001])); + fn test_multivec_distance_hamming() { + let vectors = + multivecs_of::(vec![vec![0b0000_0000, 0b0000_1111], vec![0b0000_0011]], 1); + let query: Arc = Arc::new(UInt8Array::from(vec![0b0000_0000_u8, 0b0000_1111])); let dists = multivec_distance(query.as_ref(), &vectors, DistanceType::Hamming).unwrap(); - assert_eq!(dists.len(), 1); - // One differing bit between the query and the single stored sub-vector. - assert_eq!(dists[0], 1.0 - 1.0); + + assert_eq!(dists, vec![0.0, 4.0]); + } + + #[rstest::rstest] + #[case::l2_perfect( + DistanceType::L2, + vec![1.0, 0.0, 0.0, 1.0], + vec![1.0, 0.0, 0.0, 1.0], + 0.0 + )] + #[case::cosine_perfect( + DistanceType::Cosine, + vec![1.0, 0.0, 0.0, 1.0], + vec![1.0, 0.0, 0.0, 1.0], + 0.0 + )] + #[case::dot_perfect( + DistanceType::Dot, + vec![1.0, 0.0, 0.0, 1.0], + vec![1.0, 0.0, 0.0, 1.0], + 0.0 + )] + #[case::cosine_repeated_query( + DistanceType::Cosine, + vec![0.6, 0.8], + vec![1.0, 0.0, 1.0, 0.0], + 0.8 + )] + #[case::cosine_single_query( + DistanceType::Cosine, + vec![0.0, 1.0], + vec![1.0, 0.0], + 1.0 + )] + fn test_multivec_distance_float( + #[case] distance_type: DistanceType, + #[case] vectors: Vec, + #[case] query: Vec, + #[case] expected: f32, + ) { + let vectors = multivec_of::(vectors, 2); + let query: Arc = Arc::new(Float32Array::from(query)); + + let dists = multivec_distance(query.as_ref(), &vectors, distance_type).unwrap(); + + assert!((dists[0] - expected).abs() < 1e-6); } #[test]