sqlite/ext/vec1/test/vec1pq.test

# 2026-01-17
#
# The author disclaims copyright to this source code.  In place of
# a legal notice, here is a blessing:
#
#    May you do good and not evil.
#    May you find forgiveness for yourself and forgive others.
#    May you share freely, never taking more than you give.
#
#***********************************************************************
#
#

source [file join [file dirname [info script]] vec1_common.tcl]
set testprefix vec1pq
vec1_load db

do_execsql_test 1.0 {
  CREATE VIRTUAL TABLE v1 USING vec1(v, a1);
}

do_execsql_test 1.1 {
  CREATE TABLE t1(x);
  WITH s(i) AS (
    SELECT 1 UNION ALL SELECT i+1 FROM s WHERE i<1000
  )
  INSERT INTO t1 SELECT random_vector(24) FROM s;
}

do_execsql_test 1.2 {
 INSERT INTO v1(cmd, arg) VALUES('rebuild', (
   SELECT vec1_train(x, '{codesize:8}') FROM t1
 ));
}

do_test 1.3 {
  foreach {rowid v} [db eval {SELECT rowid, x FROM t1}] {
    db eval { INSERT INTO v1 VALUES($v, $rowid) }
  }

  execsql { PRAGMA integrity_check }
} {ok}

do_execsql_test 1.4 {
  SELECT rowid, a1 FROM v1(random_vector(24), '{K:5}') WHERE a1 = 555
} {555 555}

do_execsql_test 1.5 {
  DELETE FROM v1 WHERE rowid=444;
  SELECT rowid, a1 FROM v1(random_vector(24), '{K:5}') WHERE a1 = 444;
}

do_execsql_test 1.6 {
  BEGIN;
    INSERT INTO v1(rowid, v) VALUES(10_000_000_000, random_vector(24));
    INSERT INTO v1(rowid, v) VALUES(0, random_vector(24));
    INSERT INTO v1(rowid, v) VALUES(15_000_000_000, random_vector(24));
  COMMIT;
  PRAGMA integrity_check;
} {ok}

do_execsql_test 2.0 {
  CREATE VIRTUAL TABLE v2 USING vec1(v, a2);
}

do_execsql_test 2.1 {
  WITH s(i) AS (
    SELECT 1 UNION ALL SELECT i+1 FROM s WHERE i<1000
  )
  INSERT INTO v2(rowid, v, a2) SELECT 
      i*10_000_000_000, random_vector(27), i*10_000_000_000 
  FROM s;
}

do_execsql_test 2.2 {
 INSERT INTO v2(cmd, arg) VALUES('rebuild', (
   SELECT vec1_train(v, '{codesize:9,nbucket:2,residual:0}') FROM v2
 ));
 SELECT vec1_config('nprobe', 2);
} {2.0}

do_execsql_test 2.3 {
  SELECT rowid, a2 FROM v2(random_vector(27), '{K:5}') 
  WHERE a2 = 5_550_000_000_000
} {5550000000000 5550000000000}

do_execsql_test 2.4 {
  DELETE FROM v2 WHERE rowid=5550000000000;
  SELECT rowid, a2 FROM v2(random_vector(27), '{K:5}') 
  WHERE a2 = 5_550_000_000_000
} {}

#-------------------------------------------------------------------------
reset_db
vec1_load db

do_execsql_test 3.1 {
  CREATE VIRTUAL TABLE v1 USING vec1(x);
  WITH s(i) AS (
    SELECT 1 UNION ALL SELECT i+1 FROM s WHERE i<1000
  )
  INSERT INTO v1 SELECT random_vector(32) FROM s;
}

do_execsql_test 3.2 {
  SELECT length( vec1_train(x, '{codesize:11}') ) FROM v1;
} {33816}

do_execsql_test 3.3 {
  INSERT INTO v1(cmd, arg) VALUES('rebuild', (
        SELECT vec1_train(x, '{codesize:11}') FROM v1
  ));
} {}

do_execsql_test 3.4 {
  PRAGMA integrity_check
} {ok}

#-------------------------------------------------------------------------
reset_db
vec1_load db

do_execsql_test 4.0 {
  CREATE VIRTUAL TABLE docs USING vec1(e);
  INSERT INTO docs(rowid, e) VALUES(1, random_vector(200, 10000.0));
  WITH RECURSIVE r(i) AS (
    SELECT 1 UNION ALL SELECT i+1 FROM r WHERE i<600
  ) 
  INSERT INTO docs(rowid, e) SELECT i+1, random_vector(200) FROM r;
}

do_execsql_test 4.1 {
  INSERT INTO docs(cmd, arg) VALUES('rebuild', (
    SELECT vec1_train(e, '{nbucket:2, quantizer:"pq"}') FROM docs
  ));
}

do_execsql_test 4.2 {
  PRAGMA integrity_check
} {ok}

finish_test