1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

//! In-memory graph representations.

use std::sync::Arc;

use super::storage::{DistCalculator, VectorStorage};
use arrow_array::types::Float32Type;
use lance_linalg::{distance::MetricType, MatrixView};

/// All data are stored in memory
pub struct InMemoryVectorStorage {
    row_ids: Vec<u64>,
    vectors: Arc<MatrixView<Float32Type>>,
    metric_type: MetricType,
}

impl InMemoryVectorStorage {
    pub fn new(vectors: Arc<MatrixView<Float32Type>>, metric_type: MetricType) -> Self {
        let row_ids = (0..vectors.num_rows() as u64).collect();
        Self {
            row_ids,
            vectors,
            metric_type,
        }
    }

    pub fn vector(&self, id: u32) -> &[f32] {
        self.vectors.row(id as usize).unwrap()
    }

    /// Distance between two vectors.
    pub fn distance_between(&self, a: u32, b: u32) -> f32 {
        let vector1 = self.vectors.row(a as usize).unwrap();
        let vector2 = self.vectors.row(b as usize).unwrap();
        self.metric_type.func()(vector1, vector2)
    }
}

impl VectorStorage for InMemoryVectorStorage {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn len(&self) -> usize {
        self.vectors.num_rows()
    }

    fn row_ids(&self) -> &[u64] {
        &self.row_ids
    }

    fn metric_type(&self) -> MetricType {
        self.metric_type
    }

    fn dist_calculator(&self, query: &[f32]) -> Box<dyn DistCalculator> {
        Box::new(InMemoryDistanceCal {
            vectors: self.vectors.clone(),
            query: query.to_vec(),
            metric_type: self.metric_type,
        })
    }
}

struct InMemoryDistanceCal {
    vectors: Arc<MatrixView<Float32Type>>,
    query: Vec<f32>,
    metric_type: MetricType,
}

impl DistCalculator for InMemoryDistanceCal {
    fn distance(&self, ids: &[u32]) -> Vec<f32> {
        ids.iter()
            .map(|id| {
                let vector = self.vectors.row(*id as usize).unwrap();
                self.metric_type.func()(&self.query, vector)
            })
            .collect()
    }
}