chore: fix cosine residual calculation #2015
Conversation
eddyxu
requested review from
wjones127,
westonpace,
chebbyChefNEQ and
BubbleCal
and removed request for
wjones127 and
westonpace
rust/lance-linalg/src/kernels.rs
Outdated
| <T as ArrowPrimitiveType>::Native: Float + Sum, | ||
| { | ||
| let v = arr.as_primitive::<T>(); | ||
| Ok(Arc::new(PrimitiveArray::<T>::from_iter_values(normalize(v.values()))) as ArrayRef) |
There was a problem hiding this comment.
FYI I plan to fix this soon, but normalize doesn't dispatch to the optimized SIMD kernels we have yet.
rust/lance-index/src/vector/ivf.rs
Outdated
| let num_rows = data.len() / dimension; | ||
| let num_rows = data.num_rows(); | ||
|
|
||
| let (data, metric_type) = if self.metric_type == MetricType::Cosine { |
There was a problem hiding this comment.
should we move this higher? maybe to search?
There was a problem hiding this comment.
THe shuffler is still using this piece separately, tho
rust/lance-index/src/vector/pq.rs
Outdated
| .collect::<Vec<_>>(); | ||
| let data = T::ArrayType::from(values); | ||
| FixedSizeListArray::try_new_from_values(data, self.dimension as i32)? | ||
| normalize_fsl(&fsl)? |
There was a problem hiding this comment.
Should we remove cosine type PQ entirely? And let IVF handle normalization and resid?
There was a problem hiding this comment.
^^ can we add a panic! that make sure we don't create a cosine PQ?
eddyxu
force-pushed
the
lei/pq_cosine_idx
branch
from
7ef3c0c to
9d8a2e5
Compare
eddyxu
force-pushed
the
lei/pq_cosine_idx
branch
from
9d8a2e5 to
2be5e30
Compare
rust/lance-index/src/vector/ivf.rs
Outdated
|
|
||
| // TODO: add range filter | ||
| let ivf_transform = Arc::new(IvfTransformer::new( | ||
| centroids.clone(), | ||
| metric_type, | ||
| MetricType::L2, |
There was a problem hiding this comment.
what about dot product?
| @@ -402,7 +394,7 @@ where | |||
| let mut best_kmeans = Self::empty(k, dimension, params.metric_type); | |||
| let mut best_stddev = f32::MAX; | |||
|
|
|||
| let rng = rand::rngs::SmallRng::from_entropy(); | |||
| let rng = SmallRng::from_entropy(); | |||
There was a problem hiding this comment.
nit: add a todo to use seeds
rust/lance-linalg/src/kmeans.rs
Outdated
| .collect::<Vec<_>>(); | ||
| return compute_partitions_l2(centroids_array, &normalized, dimension) | ||
| .collect(); | ||
| panic!("KMeans: should not use cosine distance to train kmeans, use L2 instead."); |
There was a problem hiding this comment.
add this check in the constructor?
There was a problem hiding this comment.
The match arm needs a branch to handle MetricType::Cosine tho.
| let partition_ids = | ||
| self.ivf | ||
| .find_partitions(&query.key, query.nprobes, self.metric_type)?; | ||
| let mut query = query.clone(); |
There was a problem hiding this comment.
how cheap is this copy, can we avoid it for the L2 path?
There was a problem hiding this comment.
This is very cheap, the query vector is zero copy (copied by pointer), other than that, there are just 6 other string/int fields.
| // vector_column, | ||
| // ))); | ||
| // }; | ||
| let mt = if metric_type == MetricType::Cosine { |
There was a problem hiding this comment.
nit: make a macro for this?
rust/lance/src/index/vector/pq.rs
Outdated
| @@ -163,6 +167,7 @@ impl VectorIndex for PQIndex { | |||
| }); | |||
| } | |||
| pre_filter.wait_for_ready().await?; | |||
| println!("PQIndex::search: metric type: {:?}", self.metric_type); | |||
rust/lance/src/index/vector/pq.rs
Outdated
| training_data = normalize_fsl(&training_data)?; | ||
| } | ||
|
|
||
| println!("PQ Training, ivf: {:?} ", ivf); |
| span!(Level::INFO, "compute residual for PQ training") | ||
| .in_scope(|| ivf2.compute_residual(&training_data, None)) | ||
| .await? |
There was a problem hiding this comment.
I think we still need to make sure we don't residulize for dot.
So, maybe something like pq_params.should_residulize?
There was a problem hiding this comment.
So the design lets the caller decide whether to run residual or not. If provided IVF, it runs residual PQ.
So PQ is orthogonal to distance type as much as possible.
There was a problem hiding this comment.
but we always residualize currently.
There was a problem hiding this comment.
need a flag for dot distance to not residualize
There was a problem hiding this comment.
Ok, checked it on the caller
No description provided.