adjust sample reader to use a LCG mapping

Author: benwtrent

server/src/main/java/org/elasticsearch/index/codec/vectors/SampleReader.java

@@ -23,6 +23,7 @@
 import org.apache.lucene.index.FloatVectorValues;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.util.Bits;
+import org.apache.lucene.util.MathUtil;
 
 import java.io.IOException;
 import java.util.Random;
@@ -81,30 +82,52 @@ public Bits getAcceptOrds(Bits acceptDocs) {
 
     static SampleReader createSampleReader(FloatVectorValues origin, int k, long seed) {
         // TODO can we do something algorithmically that aligns an ordinal with a unique integer between 0 and numVectors?
-        int[] samples = reservoirSample(origin.size(), k, seed);
-        return new SampleReader(origin, samples.length, i -> samples[i]);
+        if (k >= origin.size()) {
+            new SampleReader(origin, origin.size(), i -> i);
+        }
+        Random rnd = new Random(seed);
+        RandomLinearCongruentialMapper mapper = new RandomLinearCongruentialMapper(k, origin.size(), rnd);
+        return new SampleReader(origin, k, i -> (int) mapper.map(i));
     }
 
     /**
-     * Sample k elements from n elements according to reservoir sampling algorithm.
-     *
-     * @param n number of elements
-     * @param k number of samples
-     * @param seed random seed
-     * @return array of k samples
+     * RandomLinearCongruentialMapper is used to map a range of integers [1, n] to a range of integers [1, m]
      */
-    public static int[] reservoirSample(int n, int k, long seed) {
-        Random rnd = new Random(seed);
-        int[] reservoir = new int[k];
-        for (int i = 0; i < k; i++) {
-            reservoir[i] = i;
+    static class RandomLinearCongruentialMapper {
+        private final long n;
+        private final long m;
+        private final long multiplier;
+        private final int randomLinearShift;
+
+        RandomLinearCongruentialMapper(long smaller, long larger, Random random) {
+            assert smaller > 0 && larger > 0;
+            assert smaller < larger;
+            this.n = smaller;
+            this.m = larger;
+            this.multiplier = findLargeOddCoprime(n);
+            this.randomLinearShift = random.nextInt(0, 1024 * 1024);
+        }
+
+        // need to ensure positive modulus only
+        private static long mod(long x, long m) {
+            long r = x % m;
+            return r < 0 ? r + m : r;
         }
-        for (int i = k; i < n; i++) {
-            int j = rnd.nextInt(i + 1);
-            if (j < k) {
-                reservoir[j] = i;
+
+        long map(long i) {
+            if (i < 0 || i >= n) {
+                throw new IllegalArgumentException("i out of range");
+            }
+            long permuted = mod((i * multiplier + randomLinearShift), n);
+            return 1 + mod(permuted, m);
+        }
+
+        private static long findLargeOddCoprime(long n) {
+            long candidate = n | 1; // make sure it's odd
+            while (MathUtil.gcd(candidate, n) != 1) {
+                candidate += 2;
             }
+            return candidate;
         }
-        return reservoir;
     }
 }

server/src/test/java/org/elasticsearch/index/codec/vectors/SampleReaderTests.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.index.codec.vectors;
+
+import org.apache.lucene.util.FixedBitSet;
+import org.elasticsearch.test.ESTestCase;
+
+public class SampleReaderTests extends ESTestCase {
+
+    public void testRandomSampling() {
+        int randomLongLower = randomIntBetween(0, 1024 * 10);
+        int randomLongUpper = randomIntBetween(randomLongLower, 1024 * 100);
+        SampleReader.RandomLinearCongruentialMapper mapper = new SampleReader.RandomLinearCongruentialMapper(
+            randomLongLower,
+            randomLongUpper,
+            random()
+        );
+        FixedBitSet valueSeen = new FixedBitSet(randomLongUpper + 1);
+        for (int i = 0; i < randomLongLower; i++) {
+            long mapped = mapper.map(i);
+            assertTrue(mapped >= 0);
+            assertTrue(mapped <= randomLongUpper);
+            assertFalse(valueSeen.getAndSet((int) mapped));
+        }
+    }
+
+}