Vector Storage f16

[charles-r-earp]

/claim #3333

Based on feedback from #4032, reduce the scope of the changes to just vector_storage, not change data_types::vectors::VectorElementType.

All changes within segment crate. Changes are not visible outside segment, qdrant builds and passes tests with segment/f16 feature.

Key Changes

f16 feature (adds half dependency).

Implement data_types::primitive::PrimitiveVectorElement for f16.

In vector_storage::vector_storage_base:

#[cfg(not(feature = "f16"))]
pub(crate) type VectorStorageElementType = VectorElementType;
#[cfg(feature = "f16")]
pub(crate) type VectorStorageElementType = half::f16;

This is the type used in VectorStorageEnum:

pub enum VectorStorageEnum {
    DenseSimple(SimpleDenseVectorStorage<VectorStorageElementType>),
    DenseMemmap(Box<MemmapDenseVectorStorage<VectorStorageElementType>>),
    DenseAppendableMemmap(Box<AppendableMmapDenseVectorStorage<VectorStorageElementType>>),
    SparseSimple(SimpleSparseVectorStorage),
    MultiDenseSimple(SimpleMultiDenseVectorStorage),
}

Add new mod spaces::metric_f16 (feature f16). Implement spaces::Metric for all metrics.

For types::Distance, vector functions are generic:

impl Distance {
    pub fn preprocess_vector<T: PrimitiveVectorElement>(
            &self,
            vector: DenseVector,
        ) -> TypedDenseVector<T>
        where
            CosineMetric: Metric<T>,
            EuclidMetric: Metric<T>,
            DotProductMetric: Metric<T>,
            ManhattanMetric: Metric<T>;
    pub fn similarity<T: PrimitiveVectorElement>(&self, v1: &[T], v2: &[T]) -> ScoreType
        where
            CosineMetric: Metric<T>,
            EuclidMetric: Metric<T>,
            DotProductMetric: Metric<T>,
            ManhattanMetric: Metric<T>;
}

These are called with both VectorElementType and VectorStorageElementType.

Implementers of vector_storage::query_scorer::QueryScorer and related functions are generic over TMetric: Metric (and VectorElementType as required). For example (vector_storage::raw_scorer):

fn new_multi_scorer_with_metric<
    'a,
    TMetric: Metric<VectorStorageElementType> + Metric<VectorElementType> + 'a,
    TVectorStorage: MultiVectorStorage,
>(
    query: QueryVector,
    vector_storage: &'a TVectorStorage,
    point_deleted: &'a BitSlice,
    is_stopped: &'a AtomicBool,
) -> OperationResult<Box<dyn RawScorer + 'a>>;

This is because the query is VectorElementType, while vector_storage is VectorStorageElementType.

Test Failures

cargo test -p segment --features f16

---- multivector_hnsw_test::test_single_multi_and_dense_hnsw_equivalency stdout ----
thread 'multivector_hnsw_test::test_single_multi_and_dense_hnsw_equivalency' panicked at lib/segment/tests/integration/multivector_hnsw_test.rs:180:9:
assertion `left == right` failed
  left: [[ScoredPointOffset { idx: 593, score: 0.2239838 }, ScoredPointOffset { idx: 577, score: -0.024489254 }, ScoredPointOffset { idx: 287, score: -0.45335948 }, ScoredPointOffset { idx: 895, score: -0.493725 }]]
 right: [[ScoredPointOffset { idx: 593, score: 0.2241156 }, ScoredPointOffset { idx: 577, score: -0.024583347 }, ScoredPointOffset { idx: 287, score: -0.45343193 }, ScoredPointOffset { idx: 895, score: -0.49391016 }]]
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

---- segment_tests::test_update_named_vector stdout ----
thread 'segment_tests::test_update_named_vector' panicked at lib/segment/tests/integration/segment_tests.rs:283:13:
assertion failed: (sqrt_distance(v) - 1.).abs() < 1e-5

---- sparse_discover_test::sparse_index_discover_test stdout ----
thread 'sparse_discover_test::sparse_index_discover_test' panicked at lib/segment/tests/integration/sparse_discover_test.rs:201:9:
assertion `left == right` failed
  left: [1827, 906, 2664]
 right: [1827, 906, 2516]


failures:
    multivector_hnsw_test::test_single_multi_and_dense_hnsw_equivalency
    segment_tests::test_update_named_vector
    sparse_discover_test::sparse_index_discover_test

Potentially these failures are due to rounding error with the reduced precision of f16.

Next Steps

In vector_storage::quantized::quantized_vectors::QuantizedVectors::create_impl:

// TODO: Avoid cloning here by supporting f16 in quantization.
#[cfg(feature = "f16")]
 let vectors: Vec<_> = vectors
       .map(half::slice::HalfFloatSliceExt::to_f32_vec)
       .collect();
  #[cfg(feature = "f16")]
  let vectors = vectors.iter().map(Vec::as_slice);

quantization::encoded_vectors_u8::EncodedVectorsU8::encode takes f32 slices, which requires collecting f16 slices into a vec of f32 vecs. Could support f16 or otherwise minimize unnecessary allocation / copying.

segment::spaces::metric_f16 only has fallback implementations for Metrics. Support for avx, sse, and neon intrinsics should be straightforward. This was somewhat implemented in #4032.

[generall]

Closing in favor of #4122