Skip to content

Commit 47df4a4

Browse files
committed
more tests
1 parent 2443459 commit 47df4a4

File tree

12 files changed

+273
-98
lines changed

12 files changed

+273
-98
lines changed

fdb-extensions/src/main/java/com/apple/foundationdb/async/hnsw/HNSW.java

Lines changed: 5 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -548,8 +548,8 @@ RealVector centroidRot(@Nonnull final FhtKacRotator rotator) {
548548
}
549549

550550
@Nonnull
551-
Quantizer raBitQuantizer(@Nonnull final RealVector centroidRot) {
552-
return new RaBitQuantizer(Metric.EUCLIDEAN_METRIC, centroidRot, getConfig().getRaBitQNumExBits());
551+
Quantizer raBitQuantizer() {
552+
return new RaBitQuantizer(Metric.EUCLIDEAN_METRIC, getConfig().getRaBitQNumExBits());
553553
}
554554

555555
//
@@ -596,7 +596,7 @@ public CompletableFuture<? extends List<? extends NodeReferenceAndNode<? extends
596596
final RealVector centroidRot = centroidRot(rotator);
597597
final RealVector queryVectorRot = rotator.operateTranspose(queryVector);
598598
queryVectorTrans = queryVectorRot.subtract(centroidRot);
599-
quantizer = raBitQuantizer(centroidRot);
599+
quantizer = raBitQuantizer();
600600
} else {
601601
queryVectorTrans = queryVector;
602602
quantizer = Quantizer.noOpQuantizer(Metric.EUCLIDEAN_METRIC);
@@ -1127,7 +1127,7 @@ public CompletableFuture<Void> insert(@Nonnull final Transaction transaction, @N
11271127
final RealVector centroidRot = centroidRot(rotator);
11281128
final RealVector newVectorRot = rotator.operateTranspose(newVector);
11291129
newVectorTrans = newVectorRot.subtract(centroidRot);
1130-
quantizer = raBitQuantizer(centroidRot);
1130+
quantizer = raBitQuantizer();
11311131
} else {
11321132
newVectorTrans = newVector;
11331133
quantizer = Quantizer.noOpQuantizer(Metric.EUCLIDEAN_METRIC);
@@ -1228,7 +1228,7 @@ public CompletableFuture<Void> insertBatch(@Nonnull final Transaction transactio
12281228
if (getConfig().isUseRaBitQ()) {
12291229
rotator = new FhtKacRotator(0, getConfig().getNumDimensions(), 10);
12301230
centroidRot = centroidRot(rotator);
1231-
quantizer = raBitQuantizer(centroidRot);
1231+
quantizer = raBitQuantizer();
12321232
} else {
12331233
rotator = null;
12341234
centroidRot = null;
@@ -1913,24 +1913,6 @@ private int insertionLayer() {
19131913
return (int) Math.floor(-Math.log(u) * lambda);
19141914
}
19151915

1916-
/**
1917-
* Logs a message at the INFO level, using a consumer for lazy evaluation.
1918-
* <p>
1919-
* This approach avoids the cost of constructing the log message if the INFO
1920-
* level is disabled. The provided {@link java.util.function.Consumer} will be
1921-
* executed only when {@code logger.isInfoEnabled()} returns {@code true}.
1922-
*
1923-
* @param loggerConsumer the {@link java.util.function.Consumer} that will be
1924-
* accepted if logging is enabled. It receives the
1925-
* {@code Logger} instance and must not be null.
1926-
*/
1927-
@SuppressWarnings("PMD.UnusedPrivateMethod")
1928-
private void info(@Nonnull final Consumer<Logger> loggerConsumer) {
1929-
if (logger.isInfoEnabled()) {
1930-
loggerConsumer.accept(logger);
1931-
}
1932-
}
1933-
19341916
private static class NodeReferenceWithLayer extends NodeReferenceWithVector {
19351917
private final int layer;
19361918

fdb-extensions/src/main/java/com/apple/foundationdb/async/rabitq/RaBitEstimator.java

Lines changed: 2 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -30,15 +30,11 @@
3030
public class RaBitEstimator implements Estimator {
3131
@Nonnull
3232
private final Metric metric;
33-
@Nonnull
34-
private final RealVector centroid;
3533
private final int numExBits;
3634

3735
public RaBitEstimator(@Nonnull final Metric metric,
38-
@Nonnull final RealVector centroid,
3936
final int numExBits) {
4037
this.metric = metric;
41-
this.centroid = centroid;
4238
this.numExBits = numExBits;
4339
}
4440

@@ -47,10 +43,6 @@ public Metric getMetric() {
4743
return metric;
4844
}
4945

50-
public int getNumDimensions() {
51-
return centroid.getNumDimensions();
52-
}
53-
5446
public int getNumExBits() {
5547
return numExBits;
5648
}
@@ -78,8 +70,7 @@ private double distance(@Nonnull final RealVector query, // pre-rotated query q
7870
public Result estimateDistanceAndErrorBound(@Nonnull final RealVector query, // pre-rotated query q
7971
@Nonnull final EncodedRealVector encodedVector) {
8072
final double cb = (1 << numExBits) - 0.5;
81-
final RealVector qc = query;
82-
final double gAdd = qc.dot(qc);
73+
final double gAdd = query.dot(query);
8374
final double gError = Math.sqrt(gAdd);
8475
final RealVector totalCode = new DoubleRealVector(encodedVector.getEncodedData());
8576
final RealVector xuc = totalCode.subtract(cb);
@@ -117,7 +108,7 @@ public double getErr() {
117108

118109
@Override
119110
public String toString() {
120-
return "Estimate[" + "distance=" + distance + ", err=" + err + "]";
111+
return "estimate[" + "distance=" + distance + ", err=" + err + "]";
121112
}
122113
}
123114
}

fdb-extensions/src/main/java/com/apple/foundationdb/async/rabitq/RaBitQuantizer.java

Lines changed: 8 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -40,28 +40,20 @@ public final class RaBitQuantizer implements Quantizer {
4040
0.00, 0.15, 0.20, 0.52, 0.59, 0.71, 0.75, 0.77, 0.81
4141
};
4242

43-
@Nonnull
44-
private final RealVector centroid;
4543
final int numExBits;
4644
@Nonnull
4745
private final Metric metric;
4846

4947
public RaBitQuantizer(@Nonnull final Metric metric,
50-
@Nonnull final RealVector centroid,
5148
final int numExBits) {
52-
this.centroid = centroid;
5349
this.numExBits = numExBits;
5450
this.metric = metric;
5551
}
5652

57-
public int getNumDimensions() {
58-
return centroid.getNumDimensions();
59-
}
60-
6153
@Nonnull
6254
@Override
6355
public RaBitEstimator estimator() {
64-
return new RaBitEstimator(metric, centroid, numExBits);
56+
return new RaBitEstimator(metric, numExBits);
6557
}
6658

6759
@Nonnull
@@ -70,13 +62,6 @@ public EncodedRealVector encode(@Nonnull final RealVector data) {
7062
return encodeInternal(data).getEncodedVector();
7163
}
7264

73-
/**
74-
* Port of ex_bits_code_with_factor:
75-
* - params: data & centroid (rotated)
76-
* - forms residual internally
77-
* - computes shifted signed vector here (sign(r)*(k+0.5))
78-
* - applies C++ metric-dependent formulas exactly.
79-
*/
8065
@Nonnull
8166
Result encodeInternal(@Nonnull final RealVector data) {
8267
final int dims = data.getNumDimensions();
@@ -205,16 +190,14 @@ private QuantizeExResult quantizeEx(@Nonnull final RealVector oAbs) {
205190
return new QuantizeExResult(code, t, ipNormInv);
206191
}
207192

208-
/**
209-
* Method to compute the best factor {@code t}.
210-
* @param oAbs absolute values of a (row-wise) normalized residual; length = dim; nonnegative
211-
* @return t the rescale factor that maximizes the objective
212-
*/
213193
private double bestRescaleFactor(@Nonnull final RealVector oAbs) {
194+
214195
if (numExBits < 0 || numExBits >= TIGHT_START.length) {
215196
throw new IllegalArgumentException("numExBits out of supported range");
216197
}
217198

199+
final int numDimensions = oAbs.getNumDimensions();
200+
218201
// max_o = max(oAbs)
219202
double maxO = 0.0d;
220203
for (double v : oAbs.getData()) {
@@ -232,10 +215,10 @@ private double bestRescaleFactor(@Nonnull final RealVector oAbs) {
232215
final double tStart = tEnd * TIGHT_START[numExBits];
233216

234217
// cur_o_bar[i] = floor(tStart * oAbs[i]), but stored as int
235-
final int[] curOB = new int[getNumDimensions()];
236-
double sqrDen = getNumDimensions() * 0.25; // Σ (cur^2 + cur) starts from D/4
218+
final int[] curOB = new int[numDimensions];
219+
double sqrDen = numDimensions * 0.25; // Σ (cur^2 + cur) starts from D/4
237220
double numer = 0.0;
238-
for (int i = 0; i < getNumDimensions(); i++) {
221+
for (int i = 0; i < numDimensions; i++) {
239222
int cur = (int) ((tStart * oAbs.getComponent(i)) + EPS);
240223
curOB[i] = cur;
241224
sqrDen += (double) cur * cur + cur;
@@ -246,7 +229,7 @@ private double bestRescaleFactor(@Nonnull final RealVector oAbs) {
246229
// t_i(k->k+1) = (curOB[i] + 1) / oAbs[i]
247230

248231
final PriorityQueue<Node> pq = new PriorityQueue<>(Comparator.comparingDouble(n -> n.t));
249-
for (int i = 0; i < getNumDimensions(); i++) {
232+
for (int i = 0; i < numDimensions; i++) {
250233
final double curOAbs = oAbs.getComponent(i);
251234
if (curOAbs > 0.0) {
252235
double tNext = (curOB[i] + 1) / curOAbs;

fdb-extensions/src/main/java/com/apple/foundationdb/linear/ColumnMajorRealMatrix.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ public class ColumnMajorRealMatrix implements RealMatrix {
3434
private final Supplier<Integer> hashCodeSupplier;
3535

3636
public ColumnMajorRealMatrix(@Nonnull final double[][] data) {
37+
Preconditions.checkArgument(data.length > 0);
38+
Preconditions.checkArgument(data[0].length > 0);
3739
this.data = data;
3840
this.hashCodeSupplier = Suppliers.memoize(this::valueBasedHashCode);
3941
}

fdb-extensions/src/main/java/com/apple/foundationdb/linear/FloatRealVector.java

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,10 @@ public class FloatRealVector extends AbstractRealVector {
3636
@Nonnull
3737
private final Supplier<HalfRealVector> toHalfRealVectorSupplier;
3838

39+
public FloatRealVector(@Nonnull final Float[] floatData) {
40+
this(computeDoubleData(floatData));
41+
}
42+
3943
public FloatRealVector(@Nonnull final float[] floatData) {
4044
this(computeDoubleData(floatData));
4145
}
@@ -102,6 +106,15 @@ protected byte[] computeRawData() {
102106
return vectorBytes;
103107
}
104108

109+
@Nonnull
110+
private static double[] computeDoubleData(@Nonnull Float[] floatData) {
111+
double[] result = new double[floatData.length];
112+
for (int i = 0; i < floatData.length; i++) {
113+
result[i] = floatData[i];
114+
}
115+
return result;
116+
}
117+
105118
@Nonnull
106119
private static double[] computeDoubleData(@Nonnull float[] floatData) {
107120
double[] result = new double[floatData.length];

0 commit comments

Comments
 (0)