core/benchmarks/sharded-snapshot-build-complex.js

import fs from 'node:fs';
import path from 'node:path';
import { CondensedGraph } from '../src/core/CondensedGraph.js';
import { ShardedSnapshotBuilder } from '../src/core/shards/ShardedSnapshotBuilder.js';

function parseArgs(argv) {
  const args = new Map();
  for (let i = 2; i < argv.length; i++) {
    const value = argv[i];
    if (!value.startsWith('--')) continue;
    const [key, inline] = value.slice(2).split('=');
    if (inline !== undefined) {
      args.set(key, inline);
      continue;
    }
    const next = argv[i + 1];
    if (next && !next.startsWith('--')) {
      args.set(key, next);
      i++;
    } else {
      args.set(key, true);
    }
  }
  return args;
}

function makeRng(seed = 1337) {
  let state = seed >>> 0;
  return () => {
    state = (state * 1664525 + 1013904223) >>> 0;
    return state / 0xffffffff;
  };
}

function makeZipfSampler(count, skew, seed = 1337) {
  const weights = new Float64Array(count);
  let sum = 0;
  for (let i = 1; i <= count; i++) {
    const w = 1 / Math.pow(i, skew);
    weights[i - 1] = w;
    sum += w;
  }
  const cdf = new Float64Array(count);
  let acc = 0;
  for (let i = 0; i < count; i++) {
    acc += weights[i] / sum;
    cdf[i] = acc;
  }
  let state = seed >>> 0;
  const rand = () => {
    state = (state * 1664525 + 1013904223) >>> 0;
    return state / 0xffffffff;
  };
  return () => {
    const r = rand();
    let lo = 0;
    let hi = cdf.length - 1;
    while (lo < hi) {
      const mid = (lo + hi) >> 1;
      if (r <= cdf[mid]) {
        hi = mid;
      } else {
        lo = mid + 1;
      }
    }
    return lo;
  };
}

const args = parseArgs(process.argv);
const edges = Number(args.get('edges') || 1000000);
const users = Number(args.get('users') || 100000);
const groups = Number(args.get('groups') || 5000);
const docs = Number(args.get('docs') || 500000);
const bucketSize = Number(args.get('bucket') || 4096);
const output = args.get('output') || 'tmp/shards-complex-1m';
const userSkew = Number(args.get('user-skew') || 1.2);
const groupSkew = Number(args.get('group-skew') || 1.1);
const docSkew = Number(args.get('doc-skew') || 1.15);
const directions = args.get('directions')
  ? String(args.get('directions')).split(',')
  : ['out', 'in'];
const shardMode = args.get('shard-mode') || 'range';
const componentRelations = args.get('component-relations')
  ? String(args.get('component-relations')).split(',')
  : null;

const rng = makeRng(42);

console.log('Build complex sharded snapshot');
console.log(`  edges: ${edges}`);
console.log(`  users: ${users}`);
console.log(`  groups: ${groups}`);
console.log(`  docs: ${docs}`);
console.log(`  bucket: ${bucketSize}`);
console.log(`  output: ${output}`);
console.log(`  shard mode: ${shardMode}`);
if (componentRelations) {
  console.log(`  component relations: ${componentRelations.join(',')}`);
}
console.log(`  user-skew: ${userSkew}`);
console.log(`  group-skew: ${groupSkew}`);
console.log(`  doc-skew: ${docSkew}`);

const graph = new CondensedGraph();

const memberEdges = Math.floor(edges * 0.3);
const viewerEdges = Math.floor(edges * 0.3);
const ownerEdges = Math.floor(edges * 0.1);
const riskEdges = Math.floor(edges * 0.2);
let riskLimitEdges = edges - memberEdges - viewerEdges - ownerEdges - riskEdges;
if (riskLimitEdges < 0) riskLimitEdges = 0;

const userSampler = makeZipfSampler(users, userSkew, 7);
const groupSampler = makeZipfSampler(groups, groupSkew, 11);
const docSampler = makeZipfSampler(docs, docSkew, 13);

for (let i = 0; i < memberEdges; i++) {
  const userId = userSampler();
  const groupId = groupSampler();
  graph.addEdge(`user:${userId}`, 'member', `group:${groupId}`);
}

for (let i = 0; i < viewerEdges; i++) {
  const groupId = groupSampler();
  const docId = docSampler();
  graph.addEdge(`group:${groupId}`, 'viewer', `doc:${docId}`);
}

for (let i = 0; i < ownerEdges; i++) {
  const userId = userSampler();
  const docId = docSampler();
  graph.addEdge(`user:${userId}`, 'owner', `doc:${docId}`);
}

for (let i = 0; i < riskEdges; i++) {
  const userId = userSampler();
  const docId = docSampler();
  const value = Math.min(1, Math.max(0, rng()));
  graph.addEdge(`user:${userId}`, 'risk', `doc:${docId}`, { value, possibility: 1.0, reliability: 1.0 });
}

for (let i = 0; i < riskLimitEdges; i++) {
  const docId = docSampler();
  const value = 0.5;
  graph.addEdge(`doc:${docId}`, 'risk_limit', `doc:${docId}`, { value, possibility: 1.0, reliability: 1.0 });
}

graph.finalizePerfectHash();
graph.finalizeWaveletAdjacency({ dropAdjacencyList: true });

fs.mkdirSync(output, { recursive: true });
const snapshot = graph.toBinary();
const snapshotPath = path.join(output, 'snapshot.bin');
fs.writeFileSync(snapshotPath, new Uint8Array(snapshot));
console.log(`  snapshot MB: ${(snapshot.byteLength / 1024 / 1024).toFixed(2)}`);

const builder = new ShardedSnapshotBuilder({
  bucketSize,
  includeDirections: directions,
  shardMode,
  componentRelations
});
const manifest = builder.build(graph, output);
console.log(`  shards: ${manifest.shards.length}`);
console.log(`  manifest: ${path.join(output, 'manifest.json')}`);