musq 0.0.4

Musq is an asynchronous SQLite toolkit for Rust.
Documentation 
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
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

//! Integration tests for sqlite-vec support.

#![cfg_attr(not(feature = "vec"), allow(missing_docs))]
#![cfg(feature = "vec")]

mod support;

#[cfg(test)]
mod tests {
    use musq::{FromRow, Musq, VecBit, VecF32, VecInt8};

    use crate::support::connection;

    #[tokio::test]
    async fn vec_extension_available_on_direct_connection() -> anyhow::Result<()> {
        let conn = connection().await?;
        let version: String = musq::query_scalar("SELECT vec_version()")
            .fetch_one(&conn)
            .await?;
        assert!(version.starts_with('v'));
        Ok(())
    }

    #[tokio::test]
    async fn vec_extension_available_on_pool_connection() -> anyhow::Result<()> {
        let pool = Musq::new().open_in_memory().await?;
        let version: String = musq::query_scalar("SELECT vec_version()")
            .fetch_one(&pool)
            .await?;
        assert!(version.starts_with('v'));
        let _ = pool.close().await;
        Ok(())
    }

    #[tokio::test]
    async fn vec0_knn_query_with_vecf32() -> anyhow::Result<()> {
        let conn = connection().await?;

        musq::query("CREATE VIRTUAL TABLE vec_items USING vec0(embedding float[3])")
            .execute(&conn)
            .await?;

        musq::query("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)")
            .bind(1_i64)
            .bind(VecF32(vec![1.0, 0.0, 0.0]))
            .execute(&conn)
            .await?;
        musq::query("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)")
            .bind(2_i64)
            .bind(VecF32(vec![0.0, 1.0, 0.0]))
            .execute(&conn)
            .await?;
        musq::query("INSERT INTO vec_items(rowid, embedding) VALUES (?, ?)")
            .bind(3_i64)
            .bind(VecF32(vec![0.0, 0.0, 1.0]))
            .execute(&conn)
            .await?;

        let rows: Vec<(i64, f64)> = musq::query_as(
            "SELECT rowid, distance \
             FROM vec_items \
             WHERE embedding MATCH ? \
             ORDER BY distance \
             LIMIT 2",
        )
        .bind(VecF32(vec![0.0, 1.0, 0.0]))
        .fetch_all(&conn)
        .await?;

        assert_eq!(rows.len(), 2);
        assert_eq!(rows[0].0, 2);
        assert!(rows[0].1 <= rows[1].1);

        Ok(())
    }

    #[tokio::test]
    async fn vec_distance_functions_work() -> anyhow::Result<()> {
        let conn = connection().await?;

        let l2: f64 = musq::query_scalar("SELECT vec_distance_l2(?, ?)")
            .bind(VecF32(vec![1.0, 2.0, 3.0]))
            .bind(VecF32(vec![1.0, 2.0, 3.0]))
            .fetch_one(&conn)
            .await?;
        assert_eq!(l2, 0.0);

        let cosine: f64 = musq::query_scalar("SELECT vec_distance_cosine(?, ?)")
            .bind(VecF32(vec![1.0, 2.0, 3.0]))
            .bind(VecF32(vec![1.0, 2.0, 3.0]))
            .fetch_one(&conn)
            .await?;
        assert_eq!(cosine, 0.0);

        Ok(())
    }

    #[derive(Debug, PartialEq, FromRow)]
    struct EmbeddingRow {
        embedding: VecF32,
    }

    #[tokio::test]
    async fn from_row_with_vecf32() -> anyhow::Result<()> {
        let conn = connection().await?;

        musq::query("CREATE VIRTUAL TABLE vec_items_row USING vec0(embedding float[3])")
            .execute(&conn)
            .await?;
        musq::query("INSERT INTO vec_items_row(rowid, embedding) VALUES (?, ?)")
            .bind(1_i64)
            .bind(VecF32(vec![0.25, 0.5, 0.75]))
            .execute(&conn)
            .await?;

        let row: EmbeddingRow =
            musq::query_as("SELECT embedding FROM vec_items_row WHERE rowid = ?")
                .bind(1_i64)
                .fetch_one(&conn)
                .await?;
        assert_eq!(
            row,
            EmbeddingRow {
                embedding: VecF32(vec![0.25, 0.5, 0.75]),
            }
        );

        Ok(())
    }

    #[tokio::test]
    async fn subtype_sensitive_vector_wrappers() -> anyhow::Result<()> {
        let conn = connection().await?;

        let int8_type: String = musq::query_scalar("SELECT vec_type(vec_int8(?))")
            .bind(VecInt8(vec![-128, -1, 0, 127]))
            .fetch_one(&conn)
            .await?;
        assert_eq!(int8_type, "int8");

        let bit_type: String = musq::query_scalar("SELECT vec_type(vec_bit(?))")
            .bind(VecBit(vec![0b1010_1010, 0b0101_0101]))
            .fetch_one(&conn)
            .await?;
        assert_eq!(bit_type, "bit");

        let plain_type: String = musq::query_scalar("SELECT vec_type(?)")
            .bind(VecInt8(vec![-128, -1, 0, 127]))
            .fetch_one(&conn)
            .await?;
        assert_ne!(plain_type, "int8");

        Ok(())
    }
}