diff --git a/benchmarks/queryVector.mjs b/benchmarks/queryVector.mjs new file mode 100644 index 0000000..23a7a7b --- /dev/null +++ b/benchmarks/queryVector.mjs @@ -0,0 +1,58 @@ +import astraMongoose from '../dist/index.js'; +import mongoose from 'mongoose'; + +const { driver, tableDefinitionFromSchema } = astraMongoose; + +mongoose.set('autoCreate', false); +mongoose.set('autoIndex', false); +mongoose.setDriver(driver); + +const isTable = !!process.env.IS_TABLE; + +await mongoose.connect(process.env.ASTRA_URI, { isAstra: true, isTable }); + +const vectorField = isTable ? 'vector' : '$vector'; + +const contentSchema = new mongoose.Schema({ + text: { + type: String, + required: true + }, + [vectorField]: { + type: [Number], + validate: v => v == null || v.length === 512, + default: undefined, + dimension: 512, + index: { name: 'content_vector', vector: true } + } +}, { versionKey: false }); +const ContentModel = mongoose.model('Content', contentSchema, 'content'); + +const content = await ContentModel.findOne().select({ '*': 1 }).orFail(); +const $meta = [...content[vectorField]]; +$meta[0] = 0.001; + +const start = process.hrtime.bigint(); + +const totalQueries = 2000; +const parallelism = 10; + +for (let i = 0; i < totalQueries; i += parallelism) { + const batchSize = Math.min(parallelism, totalQueries - i); + await Promise.all( + Array.from({ length: batchSize }, () => + ContentModel.find().limit(10).select({ [vectorField]: 0 }).sort({ [vectorField]: { $meta } }) + ) + ); +} + +const end = process.hrtime.bigint(); +const seconds = Number(end - start) / 1e9; + +console.log(JSON.stringify({ + queries: totalQueries, + parallelism, + seconds: +seconds.toFixed(6), + queriesPerSecond: +(totalQueries / seconds).toFixed(6), + secondsPerBatch: +(seconds / (totalQueries / parallelism)).toFixed(6) +})); diff --git a/benchmarks/writeVector.mjs b/benchmarks/writeVector.mjs new file mode 100644 index 0000000..e5eacfd --- /dev/null +++ b/benchmarks/writeVector.mjs @@ -0,0 +1,99 @@ +import astraMongoose from '../dist/index.js'; +import fs from 'node:fs/promises'; +import mongoose from 'mongoose'; + +const { driver, tableDefinitionFromSchema } = astraMongoose; + +mongoose.set('autoCreate', false); +mongoose.set('autoIndex', false); +mongoose.setDriver(driver); + +const isTable = !!process.env.IS_TABLE; + +await mongoose.connect(process.env.ASTRA_URI, { isAstra: true, isTable }); + +const vectorField = isTable ? 'vector' : '$vector'; + +const contentSchema = new mongoose.Schema({ + text: { + type: String, + required: true + }, + [vectorField]: { + type: [Number], + validate: v => v == null || v.length === 512, + default: undefined, + dimension: 512, + index: { name: 'content_vector', vector: true } + } +}, { + versionKey: false, + autoCreate: false, + collectionOptions: { vector: { dimension: 512, metric: 'cosine' } } +}); +const ContentModel = mongoose.model('Content', contentSchema, 'content'); + +const tables = await mongoose.connection.listTables(); +const collections = await mongoose.connection.listCollections(); + +console.log('Tables', tables); +console.log('Collections', collections); + +if (isTable) { + if (collections.find(collection => collection.name === 'content')) { + await mongoose.connection.dropCollection('content'); + } + await mongoose.connection.collection('content').syncTable( + tableDefinitionFromSchema(contentSchema) + ); +} else { + const hasTable = tables.find(table => table.name === 'content'); + const hasCollection = collections.find(collection => collection.name === 'content'); + + if (hasTable) { + await mongoose.connection.dropTable('content'); + } + + if (!hasCollection) { + await ContentModel.createCollection(); + } +} + +await ContentModel.deleteMany(); +if (isTable) { + await ContentModel.syncIndexes(); +} + +const moviesPath = process.env.MOVIES_JSON_PATH || './movies.json'; + +let movies; +try { + const moviesRaw = await fs.readFile(moviesPath, 'utf8'); + movies = JSON.parse(moviesRaw); +} catch (err) { + console.error(`Failed to load movies dataset from "${moviesPath}".`); + console.error('Set the MOVIES_JSON_PATH environment variable to point to a valid movies.json file.'); + console.error('Original error:', err?.message || err); + process.exit(1); +} +const batchSize = 20; +const start = process.hrtime.bigint(); + +for (let i = 0; i < movies.length; i += batchSize) { + const batch = movies.slice(i, i + batchSize).map(m => ({ + text: [m?.title, m?.plot, m?.fullplot].filter(Boolean).join('\n\n').slice(0, 5000), + [vectorField]: m.vector + })); + + await ContentModel.insertMany(batch, { ordered: false }); +} + +const end = process.hrtime.bigint(); +const seconds = Number(end - start) / 1e9; + +console.log(JSON.stringify({ + inserted: movies.length, + seconds: +seconds.toFixed(6), + docsPerSecond: +(movies.length / seconds).toFixed(6), + secondsPerBatch: +(seconds / (movies.length / batchSize)).toFixed(6) +}));