[MASTRA-1974] PG Vector Index Update (#1971)

This PR does the following:
- adds index rebuilding for ivfflat indexes in order to handle
optimization
- adds testing suite to compare different dimensions, sizes and k values
to determine recall, latency and clustering
- adds support for other pg index types

.changeset/lemon-news-press.md

@@ -0,0 +1,6 @@
+---
+'@mastra/pg': patch
+'docs': patch
+---
+
+Update PG vector to allow for multiple index types

docs/src/pages/docs/reference/rag/pg.mdx

@@ -43,9 +43,111 @@ It provides robust vector similarity search capabilities within your existing Po
       defaultValue: "cosine",
       description: "Distance metric for similarity search",
     },
+    {
+      name: "indexConfig",
+      type: "IndexConfig",
+      isOptional: true,
+      defaultValue: "{ type: 'ivfflat' }",
+      description: "Index configuration",
+    },
+    {
+      name: "defineIndex",
+      type: "boolean",
+      isOptional: true,
+      defaultValue: "true",
+      description: "Whether to define the index",
+    },
+  ]}
+/>
+
+#### IndexConfig
+
+<PropertiesTable
+  content={[
+    {
+      name: "type",
+      type: "'flat' | 'hnsw' | 'ivfflat'",
+      description: "Index type",
+      defaultValue: "ivfflat",
+      properties: [
+        {
+          type: "string",
+          parameters: [
+            {
+              name: "flat",
+              type: "flat",
+              description: "Sequential scan (no index) that performs exhaustive search.",
+            },
+            {
+              name: "ivfflat",
+              type: "ivfflat",
+              description: "Clusters vectors into lists for approximate search.",
+            },
+            {
+              name: "hnsw",
+              type: "hnsw",
+              description: "Graph-based index offering fast search times and high recall.",
+            },
+          ],
+        },
+      ],
+    },
+    {
+      name: "ivf",
+      type: "IVFConfig",
+      isOptional: true,
+      description: "IVF configuration",
+      properties: [
+        {
+          type: "object",
+          parameters: [
+            {
+              name: "lists",
+              type: "number",
+              description: "Number of lists. If not specified, automatically calculated based on dataset size. (Minimum 100, Maximum 4000)",
+              isOptional: true,
+            }
+          ],
+        },
+      ],
+    },
+    {
+      name: "hnsw",
+      type: "HNSWConfig",
+      isOptional: true,
+      description: "HNSW configuration",
+      properties: [
+        {
+          type: "object",
+          parameters: [
+            {
+              name: "m",
+              type: "number",
+              description: "Maximum number of connections per node (default: 8)",
+              isOptional: true,
+            },
+            {
+              name: "efConstruction",
+              type: "number",
+              description: "Build-time complexity (default: 32)",
+              isOptional: true,
+            },
+          ],
+        },
+      ]
+    },
   ]}
 />
 
+#### Memory Requirements
+
+HNSW indexes require significant shared memory during construction. For 100K vectors:
+- Small dimensions (64d): ~60MB with default settings
+- Medium dimensions (256d): ~180MB with default settings
+- Large dimensions (384d+): ~250MB+ with default settings
+
+Higher M values or efConstruction values will increase memory requirements significantly. Adjust your system's shared memory limits if needed.
+
 ### upsert()
 
 <PropertiesTable
@@ -116,6 +218,31 @@ It provides robust vector similarity search capabilities within your existing Po
       defaultValue: "0",
       description: "Minimum similarity score threshold",
     },
+    {
+      name: "options",
+      type: "{ ef?: number; probes?: number }",
+      isOptional: true,
+      description: "Additional options for HNSW and IVF indexes",
+      properties: [
+        {
+          type: "object",
+          parameters: [
+            {
+              name: "ef",
+              type: "number",
+              description: "HNSW search parameter",
+              isOptional: true,
+            },
+            {
+              name: "probes",
+              type: "number",
+              description: "IVF search parameter",
+              isOptional: true,
+            },
+          ],
+        },
+      ],
+    },
   ]}
 />
 
@@ -138,10 +265,17 @@ Returns an array of index names as strings.
 Returns:
 
 ```typescript copy
-interface IndexStats {
+interface PGIndexStats {
   dimension: number;
   count: number;
   metric: "cosine" | "euclidean" | "dotproduct";
+  type: "flat" | "hnsw" | "ivfflat";
+  config: {
+    m?: number;
+    efConstruction?: number;
+    lists?: number;
+    probes?: number;
+  };
 }
 ```
 
@@ -161,6 +295,56 @@ interface IndexStats {
 
 Closes the database connection pool. Should be called when done using the store.
 
+### defineIndex()
+
+<PropertiesTable
+  content={[
+    {
+      name: "indexName",
+      type: "string",
+      description: "Name of the index to define",
+    },
+    {
+      name: "metric",
+      type: "'cosine' | 'euclidean' | 'dotproduct'",
+      isOptional: true,
+      defaultValue: "cosine",
+      description: "Distance metric for similarity search",
+    },
+    {
+      name: "indexConfig",
+      type: "IndexConfig",
+      description: "Configuration for the index type and parameters",
+    },
+  ]}
+/>
+
+Defines or redefines an index with specified metric and configuration. Will drop any existing index before creating the new one.
+
+```typescript copy
+// Define HNSW index
+await pgVector.defineIndex("my_vectors", "cosine", {
+  type: "hnsw",
+  hnsw: {
+    m: 8,
+    efConstruction: 32
+  }
+});
+
+// Define IVF index
+await pgVector.defineIndex("my_vectors", "cosine", {
+  type: "ivfflat",
+  ivf: {
+    lists: 100, 
+  }
+});
+
+// Define flat index
+await pgVector.defineIndex("my_vectors", "cosine", {
+  type: "flat" 
+});
+```
+
 ## Response Types
 
 Query results are returned in this format:
@@ -188,5 +372,11 @@ try {
 }
 ```
 
+## Best Practices
+
+- Regularly evaluate your index configuration to ensure optimal performance.
+- Adjust parameters like `lists` and `m` based on dataset size and query requirements.
+- Rebuild indexes periodically to maintain efficiency, especially after significant data changes.
+
 ### Related
 - [Metadata Filters](./metadata-filters)

stores/pg/README.md

@@ -56,25 +56,27 @@ const store = new PostgresStore({
   port: 5432,
   database: 'mastra',
   user: 'postgres',
-  password: 'postgres'
+  password: 'postgres',
 });
 
 // Create a thread
 await store.saveThread({
   id: 'thread-123',
   resourceId: 'resource-456',
   title: 'My Thread',
-  metadata: { key: 'value' }
+  metadata: { key: 'value' },
 });
 
 // Add messages to thread
-await store.saveMessages([{
-  id: 'msg-789',
-  threadId: 'thread-123',
-  role: 'user',
-  type: 'text',
-  content: [{ type: 'text', text: 'Hello' }]
-}]);
+await store.saveMessages([
+  {
+    id: 'msg-789',
+    threadId: 'thread-123',
+    role: 'user',
+    type: 'text',
+    content: [{ type: 'text', text: 'Hello' }],
+  },
+]);
 
 // Query threads and messages
 const savedThread = await store.getThread('thread-123');
@@ -97,14 +99,18 @@ Connection pool settings:
 ## Features
 
 ### Vector Store Features
+
 - Vector similarity search with cosine, euclidean, and dot product metrics
 - Advanced metadata filtering with MongoDB-like query syntax
 - Minimum score threshold for queries
 - Automatic UUID generation for vectors
 - Table management (create, list, describe, delete, truncate)
-- Uses pgvector's IVFFLAT indexing with 100 lists
+- Uses pgvector's IVFFLAT indexing with 100 lists by default
+- Supports HNSW indexing with configurable parameters
+- Supports flat indexing
 
 ### Storage Features
+
 - Thread and message storage with JSON support
 - Atomic transactions for data consistency
 - Efficient batch operations

stores/pg/docker-compose.perf.yaml

@@ -0,0 +1,21 @@
+services:
+  db:
+    image: pgvector/pgvector:pg16
+    container_name: 'pg-perf-test-db'
+    ports:
+      - '5435:5432' 
+    environment:
+      POSTGRES_USER: ${POSTGRES_USER:-postgres}
+      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-postgres}
+      POSTGRES_DB: ${POSTGRES_DB:-mastra}
+    shm_size: 1gb
+    command:
+      - "postgres"
+      - "-c"
+      - "shared_buffers=512MB"
+      - "-c"
+      - "maintenance_work_mem=1024MB"
+      - "-c"
+      - "work_mem=512MB"
+    tmpfs:
+      - /var/lib/postgresql/data 

stores/pg/package.json

@@ -19,7 +19,10 @@
     "build:watch": "pnpm build --watch",
     "pretest": "docker compose up -d && (for i in $(seq 1 30); do docker compose exec -T db pg_isready -U postgres && break || (sleep 1; [ $i -eq 30 ] && exit 1); done)",
     "test": "vitest run",
+    "pretest:perf": "docker compose -f docker-compose.perf.yaml up -d && (for i in $(seq 1 30); do docker compose -f docker-compose.perf.yaml exec -T db pg_isready -U postgres && break || (sleep 1; [ $i -eq 30 ] && exit 1); done)",
+    "test:perf": "NODE_OPTIONS='--max-old-space-size=16384' vitest run -c vitest.perf.config.ts",
     "posttest": "docker compose down -v",
+    "posttest:perf": "docker compose -f docker-compose.perf.yaml down -v",
     "pretest:watch": "docker compose up -d",
     "test:watch": "vitest watch",
     "posttest:watch": "docker compose down -v"