Skip to content

fix(linalg): correct multivector distance aggregation #8710

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yanghua merged 2 commits into from
Aug 30, 2026
Merged

fix(linalg): correct multivector distance aggregation #8710

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
33ec4e7
fix(linalg): correct multivector distance aggregation
ddupg Aug 24, 2026
eb48602
refactor: share multivec distance aggregation
ddupg Aug 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions python/python/tests/test_vector_index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
143 changes: 93 additions & 50 deletions rust/lance-linalg/src/distance.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -424,47 +429,37 @@ pub fn multivec_distance(
continue;
}

let sim = match distance_type {
DistanceType::Hamming => {
let query = query.as_primitive::<UInt8Type>().values();
query
.chunks_exact(dim)
.map(|q| {
multivector
.values()
.as_primitive::<UInt8Type>()
.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::<UInt8Type>(
query,
multivector,
dim,
hamming::hamming,
),
_ => match query.data_type() {
DataType::Float16 => multivec_distance_impl::<Float16Type>(
query,
multivector,
dim,
distance_type,
distance_type.func(),
),
DataType::Float32 => multivec_distance_impl::<Float32Type>(
query,
multivector,
dim,
distance_type,
distance_type.func(),
),
DataType::Float64 => multivec_distance_impl::<Float64Type>(
query,
multivector,
dim,
distance_type,
distance_type.func(),
),
_ => unreachable!("missed to check query type"),
},
};

dists.push(1.0 - sim);
dists.push(distance);
}
}
}
Expand All @@ -475,11 +470,8 @@ fn multivec_distance_impl<T: ArrowPrimitiveType>(
query: &dyn Array,
multivector: &FixedSizeListArray,
dim: usize,
distance_type: DistanceType,
) -> f32
where
T::Native: L2 + Cosine + Dot,
{
distance_func: DistanceFunc<T::Native>,
) -> f32 {
let query = query.as_primitive::<T>().values();
query
.chunks_exact(dim)
Expand All @@ -489,8 +481,8 @@ where
.as_primitive::<T>()
.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()
Expand All @@ -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;

Expand All @@ -529,21 +521,34 @@ mod tests {
.expect("write x86 runtime feature report");
}

/// Build a single-row `List<FixedSizeList<T, dim>>` holding one sub-vector.
fn multivec_of<T: ArrowPrimitiveType>(values: Vec<T::Native>, dim: i32) -> ListArray {
let inner = PrimitiveArray::<T>::from_iter_values(values);
/// Build `List<FixedSizeList<T, dim>>` rows from flattened sub-vector values.
fn multivecs_of<T: ArrowPrimitiveType>(rows: Vec<Vec<T::Native>>, dim: i32) -> ListArray {
let lengths = rows
.iter()
.map(|row| {
assert_eq!(row.len() % dim as usize, 0);
row.len() / dim as usize
})
.collect::<Vec<_>>();
let values = ScalarBuffer::from(rows.into_iter().flatten().collect::<Vec<_>>());
let inner = PrimitiveArray::<T>::new(values, None);
let fsl = FixedSizeListArray::try_new(
Arc::new(Field::new("item", T::DATA_TYPE, true)),
dim,
Arc::new(inner),
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<FixedSizeList<T, dim>>` row.
fn multivec_of<T: ArrowPrimitiveType>(values: Vec<T::Native>, dim: i32) -> ListArray {
multivecs_of::<T>(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
Expand Down Expand Up @@ -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::<Float32Type>(vec![1.0, 2.0], 2);
Expand Down Expand Up @@ -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::<UInt8Type>(vec![0b0000_1111, 0b0000_0000], 2);
let query: Arc<dyn Array> = Arc::new(UInt8Array::from(vec![0b0000_1111_u8, 0b0000_0001]));
fn test_multivec_distance_hamming() {
let vectors =
multivecs_of::<UInt8Type>(vec![vec![0b0000_0000, 0b0000_1111], vec![0b0000_0011]], 1);
let query: Arc<dyn Array> = 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<f32>,
#[case] query: Vec<f32>,
#[case] expected: f32,
) {
let vectors = multivec_of::<Float32Type>(vectors, 2);
let query: Arc<dyn Array> = Arc::new(Float32Array::from(query));

let dists = multivec_distance(query.as_ref(), &vectors, distance_type).unwrap();

assert!((dists[0] - expected).abs() < 1e-6);
}

#[test]
Expand Down
Loading