@mastra/pg 0.1.4 → 0.1.5-alpha.1

package/src/vector/performance.helpers.ts

@@ -0,0 +1,286 @@
+import { type IndexConfig, type IndexType } from './types';
+
+import { PgVector } from '.';
+
+export interface TestResult {
+  distribution: string;
+  dimension: number;
+  type: IndexType;
+  size: number;
+  k?: number;
+  metrics: {
+    recall?: number;
+    minRecall?: number;
+    maxRecall?: number;
+    latency?: {
+      p50: number;
+      p95: number;
+      lists?: number;
+      vectorsPerList?: number;
+      m?: number;
+      ef?: number;
+    };
+    clustering?: {
+      numLists?: number;
+      avgVectorsPerList?: number;
+      recommendedLists?: number;
+      distribution?: string;
+    };
+  };
+}
+
+export const generateRandomVectors = (count: number, dim: number) => {
+  return Array.from({ length: count }, () => {
+    return Array.from({ length: dim }, () => Math.random() * 2 - 1);
+  });
+};
+
+export const generateClusteredVectors = (count: number, dim: number, numClusters: number = 10) => {
+  // Generate cluster centers
+  const centers = Array.from({ length: numClusters }, () => Array.from({ length: dim }, () => Math.random() * 2 - 1));
+
+  // Generate vectors around centers with varying spread
+  return Array.from({ length: count }, () => {
+    // Pick a random cluster, with some clusters being more popular
+    const centerIdx = Math.floor(Math.pow(Math.random(), 2) * numClusters);
+    const center = centers[centerIdx] as number[];
+
+    // Add noise, with some vectors being further from centers
+    const spread = Math.random() < 0.8 ? 0.1 : 0.5; // 80% close, 20% far
+    return center.map(c => c + (Math.random() * spread - spread / 2));
+  });
+};
+
+// Or even more extreme:
+export const generateSkewedVectors = (count: number, dim: number) => {
+  // Create dense clusters with sparse regions
+  const vectors: number[][] = [];
+
+  const denseCount = Math.floor(count * 0.6);
+  const sparseCount = count - denseCount;
+
+  // Dense cluster (60% of vectors)
+  const denseCenter = Array.from({ length: dim }, () => Math.random() * 0.2);
+  for (let i = 0; i < denseCount; i++) {
+    vectors.push(denseCenter.map(c => c + (Math.random() * 0.1 - 0.05)));
+  }
+
+  // Scattered vectors (40%)
+  for (let i = 0; i < sparseCount; i++) {
+    vectors.push(Array.from({ length: dim }, () => Math.random() * 2 - 1));
+  }
+
+  return vectors.sort(() => Math.random() - 0.5); // Shuffle
+};
+
+export const findNearestBruteForce = (query: number[], vectors: number[][], k: number) => {
+  const similarities = vectors.map((vector, idx) => {
+    const similarity = cosineSimilarity(query, vector);
+    return { idx, dist: similarity };
+  });
+
+  const sorted = similarities.sort((a, b) => b.dist - a.dist);
+  return sorted.slice(0, k).map(x => x.idx);
+};
+
+export const calculateRecall = (actual: number[], expected: number[], k: number): number => {
+  let score = 0;
+  for (let i = 0; i < k; i++) {
+    if (actual[i] === expected[i]) {
+      score += 1;
+    } else if (expected.includes(actual[i] ?? 0)) {
+      score += 0.5;
+    }
+  }
+  return score / k;
+};
+
+export function cosineSimilarity(a: number[], b: number[]): number {
+  const dotProduct = a.reduce((sum, val, i) => sum + (val ?? 0) * (b[i] ?? 0), 0);
+  const normA = Math.sqrt(a.reduce((sum, val) => sum + val * val, 0));
+  const normB = Math.sqrt(b.reduce((sum, val) => sum + val * val, 0));
+  return dotProduct / (normA * normB);
+}
+
+export const formatTable = (data: any[], columns: string[]) => {
+  const colWidths = columns.map(col =>
+    Math.max(
+      col.length,
+      ...data.map(row => {
+        const value = row[col];
+        return value === undefined || value === null ? '-'.length : value.toString().length;
+      }),
+    ),
+  );
+
+  const topBorder = '┌' + colWidths.map(w => '─'.repeat(w)).join('┬') + '┐';
+  const headerSeparator = '├' + colWidths.map(w => '─'.repeat(w)).join('┼') + '┤';
+  const bottomBorder = '└' + colWidths.map(w => '─'.repeat(w)).join('┴') + '┘';
+
+  const header = '│' + columns.map((col, i) => col.padEnd(colWidths[i] ?? 0)).join('│') + '│';
+  const rows = data.map(
+    row =>
+      '│' +
+      columns
+        .map((col, i) => {
+          const value = row[col];
+          const displayValue = value === undefined || value === null ? '-' : value.toString();
+          return displayValue.padEnd(colWidths[i]);
+        })
+        .join('│') +
+      '│',
+  );
+
+  return [topBorder, header, headerSeparator, ...rows, bottomBorder].join('\n');
+};
+
+export const groupBy = <T, K extends keyof T>(
+  array: T[],
+  key: K | ((item: T) => string),
+  reducer?: (group: T[]) => any,
+): Record<string, any> => {
+  const grouped = array.reduce(
+    (acc, item) => {
+      const value = typeof key === 'function' ? key(item) : item[key];
+      if (!acc[value as any]) acc[value as any] = [];
+      acc[value as any]?.push(item);
+      return acc;
+    },
+    {} as Record<string, T[]>,
+  );
+
+  if (reducer) {
+    return Object.fromEntries(Object.entries(grouped).map(([key, group]) => [key, reducer(group)]));
+  }
+
+  return grouped;
+};
+
+export const calculateTimeout = (dimension: number, size: number, k: number) => {
+  let timeout = 600000;
+  if (dimension >= 1024) timeout *= 3;
+  else if (dimension >= 384) timeout *= 1.5;
+  if (size >= 10000) timeout *= 2;
+  if (k >= 75) timeout *= 1.5;
+  return timeout * 5;
+};
+
+export const baseTestConfigs = {
+  smokeTests: [{ dimension: 384, size: 1_000, k: 10, queryCount: 10 }],
+  '64': [
+    { dimension: 64, size: 100, k: 10, queryCount: 30 },
+    { dimension: 64, size: 100, k: 25, queryCount: 30 },
+    { dimension: 64, size: 100, k: 50, queryCount: 30 },
+    { dimension: 64, size: 100, k: 100, queryCount: 30 },
+    { dimension: 64, size: 1_000, k: 10, queryCount: 30 },
+    { dimension: 64, size: 1_000, k: 25, queryCount: 30 },
+    { dimension: 64, size: 1_000, k: 50, queryCount: 30 },
+    { dimension: 64, size: 1_000, k: 100, queryCount: 30 },
+    { dimension: 64, size: 10_000, k: 10, queryCount: 30 },
+    { dimension: 64, size: 100_000, k: 10, queryCount: 30 },
+    { dimension: 64, size: 100_000, k: 25, queryCount: 30 },
+    { dimension: 64, size: 100_000, k: 50, queryCount: 30 },
+    { dimension: 64, size: 100_000, k: 100, queryCount: 30 },
+    { dimension: 64, size: 500_000, k: 10, queryCount: 30 },
+    { dimension: 64, size: 1_000_000, k: 10, queryCount: 30 },
+  ],
+  '384': [
+    { dimension: 384, size: 100, k: 10, queryCount: 30 },
+    { dimension: 384, size: 100, k: 25, queryCount: 30 },
+    { dimension: 384, size: 100, k: 50, queryCount: 30 },
+    { dimension: 384, size: 100, k: 100, queryCount: 30 },
+    { dimension: 384, size: 1_000, k: 10, queryCount: 30 },
+    { dimension: 384, size: 1_000, k: 25, queryCount: 30 },
+    { dimension: 384, size: 1_000, k: 50, queryCount: 30 },
+    { dimension: 384, size: 1_000, k: 100, queryCount: 30 },
+    { dimension: 384, size: 10_000, k: 10, queryCount: 30 },
+    { dimension: 384, size: 100_000, k: 10, queryCount: 30 },
+    { dimension: 384, size: 100_000, k: 25, queryCount: 30 },
+    { dimension: 384, size: 100_000, k: 50, queryCount: 30 },
+    { dimension: 384, size: 100_000, k: 100, queryCount: 30 },
+    { dimension: 384, size: 500_000, k: 10, queryCount: 30 },
+  ],
+  '1024': [
+    { dimension: 1024, size: 100, k: 10, queryCount: 30 },
+    { dimension: 1024, size: 100, k: 25, queryCount: 30 },
+    { dimension: 1024, size: 100, k: 50, queryCount: 30 },
+    { dimension: 1024, size: 100, k: 100, queryCount: 30 },
+    { dimension: 1024, size: 1_000, k: 10, queryCount: 30 },
+    { dimension: 1024, size: 1_000, k: 25, queryCount: 30 },
+    { dimension: 1024, size: 1_000, k: 50, queryCount: 30 },
+    { dimension: 1024, size: 1_000, k: 100, queryCount: 30 },
+    { dimension: 1024, size: 10_000, k: 10, queryCount: 30 },
+    { dimension: 1024, size: 10_000, k: 25, queryCount: 30 },
+    { dimension: 1024, size: 10_000, k: 50, queryCount: 30 },
+    { dimension: 1024, size: 10_000, k: 100, queryCount: 30 },
+    { dimension: 1024, size: 50_000, k: 10, queryCount: 30 },
+    { dimension: 1024, size: 50_000, k: 25, queryCount: 30 },
+  ],
+  stressTests: [
+    // Maximum load
+    { dimension: 512, size: 1_000_000, k: 50, queryCount: 5 },
+
+    // Dense search
+    { dimension: 256, size: 1_000_000, k: 100, queryCount: 5 },
+
+    { dimension: 1024, size: 500_000, k: 50, queryCount: 5 },
+  ],
+};
+
+export interface TestConfig {
+  dimension: number;
+  size: number;
+  k: number;
+  queryCount: number;
+}
+
+export async function warmupQuery(vectorDB: PgVector, indexName: string, dimension: number, k: number) {
+  const warmupVector = generateRandomVectors(1, dimension)[0] as number[];
+  await vectorDB.query(indexName, warmupVector, k);
+}
+
+export async function measureLatency<T>(fn: () => Promise<T>): Promise<[number, T]> {
+  const start = performance.now();
+  const result = await fn();
+  const end = performance.now();
+  return [end - start, result];
+}
+
+export const getListCount = (indexConfig: IndexConfig, size: number): number | undefined => {
+  if (indexConfig.type !== 'ivfflat') return undefined;
+  if (indexConfig.ivf?.lists) return indexConfig.ivf.lists;
+  return Math.max(100, Math.min(4000, Math.floor(Math.sqrt(size) * 2)));
+};
+
+export const getHNSWConfig = (indexConfig: IndexConfig): { m: number; efConstruction: number } => {
+  return {
+    m: indexConfig.hnsw?.m ?? 8,
+    efConstruction: indexConfig.hnsw?.efConstruction ?? 32,
+  };
+};
+
+export function getSearchEf(k: number, m: number) {
+  return {
+    default: Math.max(k, m * k), // Default calculation
+    lower: Math.max(k, (m * k) / 2), // Lower quality, faster
+    higher: Math.max(k, m * k * 2), // Higher quality, slower
+  };
+}
+
+export function getIndexDescription({
+  type,
+  hnsw,
+}: {
+  type: IndexType;
+  hnsw: { m: number; efConstruction: number };
+}): string {
+  if (type === 'hnsw') {
+    return `HNSW(m=${hnsw.m},ef=${hnsw.efConstruction})`;
+  }
+
+  if (type === 'ivfflat') {
+    return `IVF`;
+  }
+
+  return 'Flat';
+}

package/src/vector/types.ts

@@ -0,0 +1,16 @@
+export type IndexType = 'ivfflat' | 'hnsw' | 'flat';
+
+interface IVFConfig {
+  lists?: number;
+}
+
+interface HNSWConfig {
+  m?: number; // Max number of connections (default: 16)
+  efConstruction?: number; // Build-time complexity (default: 64)
+}
+
+export interface IndexConfig {
+  type?: IndexType;
+  ivf?: IVFConfig;
+  hnsw?: HNSWConfig;
+}

package/src/vector/vector.performance.test.ts

@@ -0,0 +1,371 @@
+import pg from 'pg';
+import { describe, it, beforeAll, afterAll, beforeEach, afterEach } from 'vitest';
+
+import {
+  baseTestConfigs,
+  TestConfig,
+  TestResult,
+  calculateTimeout,
+  generateRandomVectors,
+  findNearestBruteForce,
+  calculateRecall,
+  formatTable,
+  groupBy,
+  measureLatency,
+  getListCount,
+  getSearchEf,
+  generateClusteredVectors,
+  generateSkewedVectors,
+  getHNSWConfig,
+  getIndexDescription,
+} from './performance.helpers';
+import { IndexConfig, IndexType } from './types';
+
+import { PgVector } from '.';
+
+interface IndexTestConfig extends IndexConfig {
+  type: IndexType;
+  rebuild?: boolean;
+}
+
+class PGPerformanceVector extends PgVector {
+  private perfPool: pg.Pool;
+
+  constructor(connectionString: string) {
+    super(connectionString);
+
+    const basePool = new pg.Pool({
+      connectionString,
+      max: 20, // Maximum number of clients in the pool
+      idleTimeoutMillis: 30000, // Close idle connections after 30 seconds
+      connectionTimeoutMillis: 2000, // Fail fast if can't connect
+    });
+
+    this.perfPool = basePool;
+  }
+
+  async bulkUpsert(indexName: string, vectors: number[][], metadata?: any[], ids?: string[]) {
+    const client = await this.perfPool.connect();
+    try {
+      await client.query('BEGIN');
+      const vectorIds = ids || vectors.map(() => crypto.randomUUID());
+
+      // Same query structure as upsert, just using unnest for bulk operation
+      const query = `
+        INSERT INTO ${indexName} (vector_id, embedding, metadata)
+        SELECT * FROM unnest(
+          $1::text[],
+          $2::vector[],
+          $3::jsonb[]
+        )
+        ON CONFLICT (vector_id)
+        DO UPDATE SET
+          embedding = EXCLUDED.embedding,
+          metadata = EXCLUDED.metadata
+        RETURNING embedding::text
+      `;
+
+      // Same parameter structure as upsert, just as arrays
+      await client.query(query, [
+        vectorIds,
+        vectors.map(v => `[${v.join(',')}]`),
+        (metadata || vectors.map(() => ({}))).map(m => JSON.stringify(m)),
+      ]);
+      await client.query('COMMIT');
+      return vectorIds;
+    } catch (error) {
+      await client.query('ROLLBACK');
+      throw error;
+    } finally {
+      client.release();
+    }
+  }
+}
+
+const warmupCache = new Map<string, boolean>();
+async function smartWarmup(
+  vectorDB: PGPerformanceVector,
+  testIndexName: string,
+  indexType: string,
+  dimension: number,
+  k: number,
+) {
+  const cacheKey = `${dimension}-${k}-${indexType}`;
+  if (!warmupCache.has(cacheKey)) {
+    console.log(`Warming up ${indexType} index for ${dimension}d vectors, k=${k}`);
+    const warmupVector = generateRandomVectors(1, dimension)[0] as number[];
+    await vectorDB.query(testIndexName, warmupVector, k);
+    warmupCache.set(cacheKey, true);
+  }
+}
+
+const connectionString = process.env.DB_URL || `postgresql://postgres:postgres@localhost:5435/mastra`;
+describe('PostgreSQL Index Performance', () => {
+  let vectorDB: PGPerformanceVector;
+  const testIndexName = 'test_index_performance';
+  const results: TestResult[] = [];
+
+  const indexConfigs: IndexTestConfig[] = [
+    { type: 'flat' }, // Test flat/linear search as baseline
+    { type: 'ivfflat', ivf: { lists: 100 } }, // Test IVF with fixed lists
+    { type: 'ivfflat', rebuild: true }, // Test IVF with calculated lists and rebuild
+    { type: 'hnsw' }, // Test HNSW with default parameters
+    { type: 'hnsw', hnsw: { m: 16, efConstruction: 64 } }, // Test HNSW with custom parameters
+  ];
+  beforeAll(async () => {
+    // Initialize PGPerformanceVector
+    vectorDB = new PGPerformanceVector(connectionString);
+  });
+  beforeEach(async () => {
+    await vectorDB.deleteIndex(testIndexName);
+  });
+
+  afterEach(async () => {
+    await vectorDB.deleteIndex(testIndexName);
+  });
+
+  afterAll(async () => {
+    await vectorDB.disconnect();
+    analyzeResults(results);
+  });
+
+  // Combine all test configs
+  const allConfigs: TestConfig[] = [
+    ...baseTestConfigs['64'],
+    ...baseTestConfigs['384'],
+    ...baseTestConfigs['1024'],
+    ...baseTestConfigs.smokeTests,
+    ...baseTestConfigs.stressTests,
+  ];
+
+  // For each index config
+  for (const indexConfig of indexConfigs) {
+    const indexType = indexConfig.type;
+    const rebuild = indexConfig.rebuild ?? false;
+    const hnswConfig = getHNSWConfig(indexConfig);
+    const indexDescription = getIndexDescription({
+      type: indexType,
+      hnsw: hnswConfig,
+    });
+
+    describe(`Index: ${indexDescription}`, () => {
+      for (const testConfig of allConfigs) {
+        const timeout = calculateTimeout(testConfig.dimension, testConfig.size, testConfig.k);
+        const testDesc = `dim=${testConfig.dimension} size=${testConfig.size} k=${testConfig.k}`;
+
+        for (const [distType, generator] of Object.entries(distributions)) {
+          it(
+            testDesc,
+            async () => {
+              const testVectors = generator(testConfig.size, testConfig.dimension);
+              const queryVectors = generator(testConfig.queryCount, testConfig.dimension);
+
+              // Create index and insert vectors
+              const lists = getListCount(indexConfig, testConfig.size);
+
+              await vectorDB.createIndex(
+                testIndexName,
+                testConfig.dimension,
+                'cosine',
+                indexConfig,
+                indexType === 'ivfflat',
+              );
+
+              console.log(
+                `Batched bulk upserting ${testVectors.length} ${distType} vectors into index ${testIndexName}`,
+              );
+              const batchSizes = splitIntoRandomBatches(testConfig.size, testConfig.dimension);
+              await batchedBulkUpsert(vectorDB, testIndexName, testVectors, batchSizes);
+              if (indexType === 'hnsw' || rebuild) {
+                console.log('rebuilding index');
+                await vectorDB.defineIndex(testIndexName, 'cosine', indexConfig);
+                console.log('index rebuilt');
+              }
+              await smartWarmup(vectorDB, testIndexName, indexType, testConfig.dimension, testConfig.k);
+
+              // For HNSW, test different EF values
+              const efValues = indexType === 'hnsw' ? getSearchEf(testConfig.k, hnswConfig.m) : { default: undefined };
+
+              for (const [efType, ef] of Object.entries(efValues)) {
+                const recalls: number[] = [];
+                const latencies: number[] = [];
+
+                for (const queryVector of queryVectors) {
+                  const expectedNeighbors = findNearestBruteForce(queryVector, testVectors, testConfig.k);
+
+                  const [latency, actualResults] = await measureLatency(async () =>
+                    vectorDB.query(
+                      testIndexName,
+                      queryVector,
+                      testConfig.k,
+                      undefined,
+                      false,
+                      0,
+                      { ef }, // For HNSW
+                    ),
+                  );
+
+                  const actualNeighbors = actualResults.map(r => r.metadata?.index);
+                  const recall = calculateRecall(actualNeighbors, expectedNeighbors, testConfig.k);
+                  recalls.push(recall);
+                  latencies.push(latency);
+                }
+
+                const sorted = [...latencies].sort((a, b) => a - b);
+                results.push({
+                  distribution: distType,
+                  dimension: testConfig.dimension,
+                  size: testConfig.size,
+                  k: testConfig.k,
+                  type: indexType,
+                  metrics: {
+                    recall: recalls.length > 0 ? recalls.reduce((a, b) => a + b, 0) / recalls.length : 0,
+                    minRecall: Math.min(...recalls),
+                    maxRecall: Math.max(...recalls),
+                    latency: {
+                      p50: sorted[Math.floor(sorted.length * 0.5)],
+                      p95: sorted[Math.floor(sorted.length * 0.95)],
+                      ...(indexType === 'ivfflat' && {
+                        lists,
+                        vectorsPerList: Math.round(testConfig.size / (lists || 1)),
+                      }),
+                      ...(indexType === 'hnsw' && {
+                        m: hnswConfig.m,
+                        efConstruction: hnswConfig.efConstruction,
+                        ef,
+                        efType,
+                      }),
+                    },
+                    ...(indexType === 'ivfflat' && {
+                      clustering: {
+                        numLists: lists,
+                        avgVectorsPerList: testConfig.size / (lists || 1),
+                        recommendedLists: Math.floor(Math.sqrt(testConfig.size)),
+                        distribution: distType,
+                      },
+                    }),
+                  },
+                });
+              }
+            },
+            timeout,
+          );
+        }
+      }
+    });
+  }
+});
+
+function analyzeResults(results: TestResult[]) {
+  const byType = groupBy(results, (r: TestResult) => r.type);
+  Object.entries(byType).forEach(([type, typeResults]) => {
+    console.log(`\n=== ${type.toUpperCase()} Index Analysis ===\n`);
+
+    const byDimension = groupBy(typeResults, (r: TestResult) => r.dimension.toString());
+    Object.entries(byDimension).forEach(([dim, dimResults]) => {
+      console.log(`\n--- Analysis for ${dim} dimensions ---\n`);
+
+      // Combined Performance Analysis
+      const columns = ['Distribution', 'Dataset Size', 'K'];
+      if (type === 'hnsw') {
+        columns.push('M', 'EF Construction', 'EF', 'EF Type');
+      } else if (type === 'ivfflat') {
+        columns.push('Lists', 'Vectors/List');
+      }
+      columns.push('Min Recall', 'Avg Recall', 'Max Recall', 'P50 (ms)', 'P95 (ms)');
+
+      const performanceData = Object.values(
+        groupBy(
+          dimResults,
+          (r: any) => `${r.size}-${r.k}-${type === 'ivfflat' ? r.metrics.latency.lists : r.metrics.latency.m}`,
+          (results: any[]) => {
+            const sortedResults = [...results].sort(
+              (a, b) =>
+                ['random', 'clustered', 'skewed', 'mixed'].indexOf(a.distribution) -
+                ['random', 'clustered', 'skewed', 'mixed'].indexOf(b.distribution),
+            );
+            return sortedResults.map(result => ({
+              Distribution: result.distribution,
+              'Dataset Size': result.size,
+              K: result.k,
+              ...(type === 'ivfflat'
+                ? {
+                    Lists: result.metrics.latency.lists,
+                    'Vectors/List': result.metrics.latency.vectorsPerList,
+                  }
+                : {}),
+              ...(type === 'hnsw'
+                ? {
+                    M: result.metrics.latency.m,
+                    'EF Construction': result.metrics.latency.efConstruction,
+                    EF: result.metrics.latency.ef,
+                    'EF Type': result.metrics.latency.efType,
+                  }
+                : {}),
+              'Min Recall': result.metrics.minRecall.toFixed(3),
+              'Avg Recall': result.metrics.recall.toFixed(3),
+              'Max Recall': result.metrics.maxRecall.toFixed(3),
+              'P50 (ms)': result.metrics.latency.p50.toFixed(2),
+              'P95 (ms)': result.metrics.latency.p95.toFixed(2),
+            }));
+          },
+        ),
+      ).flat();
+
+      console.log(formatTable(performanceData, columns));
+    });
+  });
+}
+
+function splitIntoRandomBatches(total: number, dimension: number): number[] {
+  const batches: number[] = [];
+  let remaining = total;
+
+  const batchRange = dimension === 1024 ? 5000 : 15000;
+
+  while (remaining > 0) {
+    const batchSize = Math.min(remaining, batchRange + Math.floor(Math.random() * batchRange));
+    batches.push(batchSize);
+    remaining -= batchSize;
+  }
+
+  return batches;
+}
+
+async function batchedBulkUpsert(
+  vectorDB: PGPerformanceVector,
+  testIndexName: string,
+  vectors: number[][],
+  batchSizes: number[],
+) {
+  let offset = 0;
+  const vectorIds = vectors.map((_, idx) => `vec_${idx}`);
+  const metadata = vectors.map((_, idx) => ({ index: idx }));
+
+  for (const size of batchSizes) {
+    const batch = vectors.slice(offset, offset + size);
+    const batchIds = vectorIds.slice(offset, offset + size);
+    const batchMetadata = metadata.slice(offset, offset + size);
+    await vectorDB.bulkUpsert(testIndexName, batch, batchMetadata, batchIds);
+    offset += size;
+    console.log(`${offset} of ${vectors.length} vectors upserted`);
+  }
+}
+
+const distributions = {
+  random: generateRandomVectors,
+  clustered: generateClusteredVectors,
+  skewed: generateSkewedVectors,
+  mixed: (size: number, dimension: number) => {
+    const generators = [generateRandomVectors, generateClusteredVectors, generateSkewedVectors];
+    const batchSizes = splitIntoRandomBatches(size, dimension);
+
+    let vectors: number[][] = [];
+    for (const batchSize of batchSizes) {
+      const generator = generators[Math.floor(Math.random() * generators.length)];
+      vectors = vectors.concat(generator(batchSize, dimension));
+    }
+
+    return vectors;
+  },
+};

package/vitest.config.ts

@@ -4,6 +4,7 @@ export default defineConfig({
   test: {
     environment: 'node',
     include: ['src/**/*.test.ts'],
+    exclude: ['src/**/*.performance.test.ts'],
     coverage: {
       reporter: ['text', 'json', 'html'],
     },

package/vitest.perf.config.ts

@@ -0,0 +1,8 @@
+import { defineConfig } from 'vitest/config';
+
+export default defineConfig({
+  test: {
+    environment: 'node',
+    include: ['src/**/*.performance.test.ts'],
+  },
+});