Improved SQLite examples

Author: ankane

README.md

@@ -306,12 +306,14 @@ class AddEmbeddingToItems < ActiveRecord::Migration[8.0]
   def change
     # Rails 8+
     create_virtual_table :items, :vec0, [
+      "id integer PRIMARY KEY AUTOINCREMENT NOT NULL",
       "embedding float[3] distance_metric=L2"
     ]
 
     # Rails < 8
     execute <<~SQL
       CREATE VIRTUAL TABLE items USING vec0(
+        id integer PRIMARY KEY AUTOINCREMENT NOT NULL,
         embedding float[3] distance_metric=L2
       )
     SQL
@@ -329,16 +331,6 @@ ActiveRecord::SchemaDumper.ignore_tables += [
 ]
 ```
 
-Create a model with `rowid` as the primary key
-
-```ruby
-class Item < ApplicationRecord
-  self.primary_key = "rowid"
-
-  has_neighbors :embedding, dimensions: 3
-end
-```
-
 Get the `k` nearest neighbors
 
 ```ruby
@@ -348,7 +340,7 @@ Item.where("embedding MATCH ?", [1, 2, 3].to_s).where(k: 5).order(:distance)
 Filter by primary key
 
 ```ruby
-Item.where(rowid: [2, 3]).where("embedding MATCH ?", [1, 2, 3].to_s).where(k: 5).order(:distance)
+Item.where(id: [2, 3]).where("embedding MATCH ?", [1, 2, 3].to_s).where(k: 5).order(:distance)
 ```
 
 ### Int8 Vectors

test/sqlite_virtual_test.rb

@@ -23,7 +23,7 @@ def test_euclidean
 
     relation = SqliteVecItem.where("embedding MATCH ?", [1, 1, 1].to_s).order(:distance).limit(3)
     assert_equal [1, 3, 2], relation.all.map(&:id)
-    assert_equal [1, 3, 2], relation.pluck(:rowid)
+    assert_equal [1, 3, 2], relation.pluck(:id)
     assert_elements_in_delta [0, 1, Math.sqrt(3)], relation.pluck(:distance)
     assert_match "SCAN vec_items VIRTUAL TABLE INDEX", relation.explain.inspect
 
@@ -51,10 +51,17 @@ def test_where_k
     assert SqliteVecItem.where.not(embedding: nil).where("embedding MATCH ? AND k = ?", "[0, 0, 0]", 3).order(:distance).load
   end
 
-  def test_where_rowid
+  def test_where_id
     create_items(SqliteVecItem, :embedding)
 
-    relation = SqliteVecItem.where(rowid: [2, 3]).where("embedding MATCH ?", [1, 1, 1].to_s).where(k: 5).order(:distance)
-    assert_equal [3, 2], relation.pluck(:rowid)
+    relation = SqliteVecItem.where(id: [2, 3]).where("embedding MATCH ?", [1, 1, 1].to_s).where(k: 5).order(:distance)
+    assert_equal [3, 2], relation.pluck(:id)
+  end
+
+  def test_create_returning_id
+    item = SqliteVecItem.create!(embedding: [1, 2, 3])
+    # TODO figure out why id not set
+    assert_nil item.id
+    assert_kind_of Integer, SqliteVecItem.last.id
   end
 end

test/support/sqlite.rb

@@ -15,23 +15,27 @@ class SqliteRecord < ActiveRecord::Base
 
   if ActiveRecord::VERSION::MAJOR >= 8
     create_virtual_table :vec_items, :vec0, [
+      "id integer PRIMARY KEY AUTOINCREMENT NOT NULL",
       "embedding float[3] distance_metric=L2"
     ]
   else
     execute <<~SQL
       CREATE VIRTUAL TABLE vec_items USING vec0(
+        id integer PRIMARY KEY AUTOINCREMENT NOT NULL,
         embedding float[3] distance_metric=L2
       )
     SQL
   end
 
   if ActiveRecord::VERSION::MAJOR >= 8
     create_virtual_table :cosine_items, :vec0, [
+      "id integer PRIMARY KEY AUTOINCREMENT NOT NULL",
       "embedding float[3] distance_metric=cosine"
     ]
   else
     execute <<~SQL
       CREATE VIRTUAL TABLE cosine_items USING vec0(
+        id integer PRIMARY KEY AUTOINCREMENT NOT NULL,
         embedding float[3] distance_metric=cosine
       )
     SQL
@@ -47,13 +51,11 @@ class SqliteItem < SqliteRecord
 
 class SqliteVecItem < SqliteRecord
   has_neighbors :embedding, dimensions: 3
-  self.primary_key = "rowid"
   self.table_name = "vec_items"
 end
 
 class SqliteCosineItem < SqliteRecord
   has_neighbors :embedding, dimensions: 3
-  self.primary_key = "rowid"
   self.table_name = "cosine_items"
 end