Quelle  test_storage_ext_vec.js

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

// This file tests support for the sqlite-vec extension.

function tensorToBlob(tensor) {
  return new Uint8ClampedArray(new Float32Array(tensor).buffer);
}

add_setup(async function () {
  cleanup();
});

add_task(async function test_synchronous() {
  info("Testing synchronous connection");
  let conn = getOpenedUnsharedDatabase();
  Assert.throws(
    () =>
      conn.executeSimpleSQL(
        `CREATE VIRTUAL TABLE test USING vec0(
          embedding FLOAT[4]
        );`
      ),
    /NS_ERROR_FAILURE/,
    "Should not be able to use vec without loading the extension"
  );

  await loadExtension(conn);

  conn.executeSimpleSQL(
    `
    CREATE VIRTUAL TABLE test USING vec0(
      embedding FLOAT[4]
    )
    `
  );

  let stmt = conn.createStatement(
    `
    INSERT INTO test(rowid, embedding)
    VALUES (1, :vector)
    `
  );
  stmt.bindBlobByName("vector", tensorToBlob([0.3, 0.3, 0.3, 0.3]));
  stmt.executeStep();
  stmt.reset();
  stmt.finalize();

  stmt = conn.createStatement(
    `
    SELECT
      rowid,
      distance
    FROM test
    WHERE embedding MATCH :vector
    ORDER BY distance
    LIMIT 1
    `
  );
  stmt.bindBlobByName("vector", tensorToBlob([0.3, 0.3, 0.3, 0.3]));
  Assert.ok(stmt.executeStep());
  Assert.equal(stmt.getInt32(0), 1);
  Assert.equal(stmt.getDouble(1), 0.0);
  stmt.reset();
  stmt.finalize();

  cleanup();
});

add_task(async function test_asynchronous() {
  info("Testing asynchronous connection");
  let conn = await openAsyncDatabase(getTestDB());

  await Assert.rejects(
    executeSimpleSQLAsync(
      conn,
      `
      CREATE VIRTUAL TABLE test USING vec0(
        embedding float[4]
      )
      `
    ),
    err => err.message.startsWith("no such module"),
    "Should not be able to use vec without loading the extension"
  );

  await loadExtension(conn);

  await executeSimpleSQLAsync(
    conn,
    `
    CREATE VIRTUAL TABLE test USING vec0(
      embedding float[4]
    )
    `
  );

  await asyncClose(conn);
  await IOUtils.remove(getTestDB().path, { ignoreAbsent: true });
});

add_task(async function test_clone() {
  info("Testing cloning synchronous connection loads extensions in clone");
  let conn1 = getOpenedUnsharedDatabase();
  await loadExtension(conn1);

  let conn2 = conn1.clone(false);
  conn2.executeSimpleSQL(
    `
    CREATE VIRTUAL TABLE test USING vec0(
      embedding float[4]
    )
    `
  );

  conn2.close();
  cleanup();
});

add_task(async function test_asyncClone() {
  info("Testing asynchronously cloning connection loads extensions in clone");
  let conn1 = getOpenedUnsharedDatabase();
  await loadExtension(conn1);

  let conn2 = await asyncClone(conn1, false);
  await executeSimpleSQLAsync(
    conn2,
    `
    CREATE VIRTUAL TABLE test USING vec0(
      embedding float[4]
    )
    `
  );

  await asyncClose(conn2);
  await asyncCleanup();
});

async function loadExtension(conn, ext = "vec") {
  await new Promise((resolve, reject) => {
    conn.loadExtension(ext, status => {
      if (Components.isSuccessCode(status)) {
        resolve();
      } else {
        reject(status);
      }
    });
  });
}

add_task(async function test_invariants() {
  // Test some invariants of the vec extension that we rely upon, so that if
  // the behavior changes we can catch it.
  let conn = getOpenedUnsharedDatabase();
  await loadExtension(conn);

  conn.executeSimpleSQL(
    `
    CREATE VIRTUAL TABLE vectors USING vec0(
      embedding FLOAT[4]
    )
    `
  );
  conn.executeSimpleSQL(
    `
    CREATE TABLE relations (
      rowid INTEGER PRIMARY KEY,
      content TEXT
    )
    `
  );

  let rowids = [];
  let insertRelStmt = conn.createStatement(
    `
    INSERT INTO relations (rowid, content)
    VALUES (NULL, "test")
    RETURNING rowid
    `
  );
  Assert.ok(insertRelStmt.executeStep());
  rowids.push(insertRelStmt.getInt32(0));
  insertRelStmt.reset();
  Assert.ok(insertRelStmt.executeStep());
  rowids.push(insertRelStmt.getInt32(0));
  insertRelStmt.reset();

  // Try to insert the same rowid twice in the vec table.
  let insertVecStmt = conn.createStatement(
    `
    INSERT INTO vectors (rowid, embedding)
    VALUES (:rowid, :vector)
    `
  );
  insertVecStmt.bindByName("rowid", rowids[0]);
  insertVecStmt.bindBlobByName("vector", tensorToBlob([0.1, 0.1, 0.1, 0.1]));
  insertVecStmt.executeStep();
  insertVecStmt.reset();

  let deleteStmt = conn.createStatement(
    `
    DELETE FROM vectors WHERE rowid = :rowid
    `
  );
  deleteStmt.bindByName("rowid", rowids[0]);
  deleteStmt.executeStep();
  deleteStmt.finalize();

  insertVecStmt.bindByName("rowid", rowids[0]);
  insertVecStmt.bindBlobByName("vector", tensorToBlob([0.2, 0.2, 0.2, 0.2]));
  insertVecStmt.executeStep();
  insertVecStmt.reset();

  let selectStmt = conn.createStatement(
    `
    SELECT
      rowid,
      vec_to_json(embedding)
    FROM vectors
    `
  );
  let count = 0;
  while (selectStmt.executeStep()) {
    count++;
    Assert.equal(selectStmt.getInt32(0), rowids[0]);
    Assert.equal(
      selectStmt.getUTF8String(1).replace(/(?<=[0-9])0+/g, ""),
      "[0.2,0.2,0.2,0.2]"
    );
  }
  Assert.equal(count, 1, "Should have one row in the vec table");
  selectStmt.reset();

  Assert.ok(insertRelStmt.executeStep());
  rowids.push(insertRelStmt.getInt32(0));
  insertRelStmt.finalize();
  insertVecStmt.bindByName("rowid", rowids[2]);
  insertVecStmt.bindBlobByName("vector", tensorToBlob([0.3, 0.3, 0.3, 0.3]));
  insertVecStmt.executeStep();
  insertVecStmt.finalize();

  let expected = [
    { rowid: rowids[0], vector: "[0.2,0.2,0.2,0.2]" },
    { rowid: rowids[2], vector: "[0.3,0.3,0.3,0.3]" },
  ];
  count = 0;
  for (let i = 0; selectStmt.executeStep(); i++) {
    count++;
    Assert.equal(selectStmt.getInt32(0), expected[i].rowid);
    Assert.equal(
      selectStmt.getUTF8String(1).replace(/(?<=[0-9])0+/g, ""),
      expected[i].vector
    );
  }
  Assert.equal(count, 2, "Should have two rows in the vec table");
  selectStmt.finalize();

  // TODO: In the future add testing for RETURNING and UPSERT as those are
  // currently broken. See:
  // https://github.com/asg017/sqlite-vec/issues/127
  // https://github.com/asg017/sqlite-vec/issues/229

  cleanup();
});