diskann-inmem 0.56.0

DiskANN3 is a composable library for bringing scalable, accurate and cost-effective vector indexing to multiple databases.
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

/*
 * Copyright (c) Microsoft Corporation.
 * Licensed under the MIT license.
 */

use diskann_utils::{io::read_bin, views::Matrix};
use half::f16;

use super::datatype::{DataType, Dataset, Preprocess, SliceMut};

pub(crate) fn load_and_convert<IO>(
    io: &mut IO,
    src: DataType,
    target: DataType,
    ops: &[Preprocess],
) -> anyhow::Result<Dataset>
where
    IO: std::io::Read + std::io::Seek,
{
    let mut data = match src {
        DataType::F32 => Dataset::from(read_bin::<f32>(io)?),
        DataType::F16 => Dataset::from(read_bin::<f16>(io)?),
        DataType::U8 => Dataset::from(read_bin::<u8>(io)?),
        DataType::I8 => Dataset::from(read_bin::<i8>(io)?),
    };

    for op in ops {
        data.preprocess(op);
    }

    if src == target {
        return Ok(data);
    }

    let dst = match target {
        DataType::F32 => {
            let mut dst = Matrix::new(0.0f32, data.nrows(), data.ncols());
            SliceMut::from(dst.as_mut_slice()).convert_lossless(data.as_slice())?;
            Dataset::from(dst)
        }
        DataType::F16 => {
            let mut dst = Matrix::new(f16::from_f32(0.0f32), data.nrows(), data.ncols());
            SliceMut::from(dst.as_mut_slice()).convert_lossless(data.as_slice())?;
            Dataset::from(dst)
        }
        DataType::U8 => {
            let mut dst = Matrix::new(0u8, data.nrows(), data.ncols());
            SliceMut::from(dst.as_mut_slice()).convert_lossless(data.as_slice())?;
            Dataset::from(dst)
        }
        DataType::I8 => {
            let mut dst = Matrix::new(0i8, data.nrows(), data.ncols());
            SliceMut::from(dst.as_mut_slice()).convert_lossless(data.as_slice())?;
            Dataset::from(dst)
        }
    };

    Ok(dst)
}