WeightedRandomGenerator Class — drizzle-orm Architecture

Source Code

drizzle-seed/src/services/Generators.ts lines 2810–2871

export class WeightedRandomGenerator extends AbstractGenerator<{ weight: number; value: AbstractGenerator<any> }[]> {
	static override readonly entityKind: string = 'WeightedRandomGenerator';

	private state: {
		rng: prand.RandomGenerator;
		weightedIndices: number[];
	} | undefined;

	override init({ count, seed }: { count: number; seed: number }) {
		const weights = this.params.map((weightedGen) => weightedGen.weight);
		const weightedIndices = getWeightedIndices(weights);

		let idx: number, valueIdx: number, tempRng = prand.xoroshiro128plus(seed);
		const indicesCounter: { [key: number]: number } = {};
		for (let i = 0; i < count; i++) {
			[idx, tempRng] = prand.uniformIntDistribution(0, weightedIndices.length - 1, tempRng);
			valueIdx = weightedIndices[idx]!;
			if (!Object.hasOwn(indicesCounter, valueIdx)) indicesCounter[valueIdx] = 0;
			indicesCounter[valueIdx]! += 1;
		}

		for (const [idx, weightedGen] of this.params.entries()) {
			weightedGen.value.isUnique = this.isUnique;
			weightedGen.value.dataType = this.dataType;
			weightedGen.value.init({ count: indicesCounter[idx]!, seed });

			if (
				weightedGen.value.uniqueVersionOfGen !== undefined
				&& weightedGen.value.isUnique === true
			) {
				const uniqueGen = new weightedGen.value.uniqueVersionOfGen({
					...weightedGen.value.params,
				});
				uniqueGen.init({
					count: indicesCounter[idx]!,
					seed,
				});
				uniqueGen.isUnique = weightedGen.value.isUnique;
				uniqueGen.dataType = weightedGen.value.dataType;

				weightedGen.value = uniqueGen;
			}
		}

		const rng = prand.xoroshiro128plus(seed);

		this.state = { weightedIndices, rng };
	}

	generate({ i }: { i: number }) {
		if (this.state === undefined) {
			throw new Error('state is not defined.');
		}

		let idx: number;
		[idx, this.state.rng] = prand.uniformIntDistribution(0, this.state.weightedIndices.length - 1, this.state.rng);
		const generatorIdx = this.state.weightedIndices[idx] as number;
		const value = this.params[generatorIdx]!.value.generate({ i });

		return value;
	}
}