diff --git a/src/main/java/ru/rt/restream/reindexer/Reindexer.java b/src/main/java/ru/rt/restream/reindexer/Reindexer.java index 6213d53..0de3798 100644 --- a/src/main/java/ru/rt/restream/reindexer/Reindexer.java +++ b/src/main/java/ru/rt/restream/reindexer/Reindexer.java @@ -282,28 +282,16 @@ public Query query(String namespaceName, Class clazz) { return new Query<>(this, namespace, null); } - /** - * ONLY FOR TEST PURPOSES! - */ - @Deprecated public void addIndex(String namespaceName, ReindexerIndex index) { IndexDefinition indexDefinition = IndexDefinition.fromIndex(index); binding.addIndex(namespaceName, indexDefinition); } - /** - * ONLY FOR TEST PURPOSES! - */ - @Deprecated public void updateIndex(String namespaceName, ReindexerIndex index) { IndexDefinition indexDefinition = IndexDefinition.fromIndex(index); binding.updateIndex(namespaceName, indexDefinition); } - /** - * ONLY FOR TEST PURPOSES! - */ - @Deprecated public void dropIndex(String namespaceName, String indexName) { binding.dropIndex(namespaceName, indexName); } diff --git a/src/main/java/ru/rt/restream/reindexer/annotations/Hnsw.java b/src/main/java/ru/rt/restream/reindexer/annotations/Hnsw.java index 82d9b74..4b21e39 100644 --- a/src/main/java/ru/rt/restream/reindexer/annotations/Hnsw.java +++ b/src/main/java/ru/rt/restream/reindexer/annotations/Hnsw.java @@ -79,4 +79,40 @@ */ boolean multithreading() default false; + /** + * Optional quantization configuration for HNSW index. + * + *

When {@code quantizationConfig.quantizationType} is empty, the whole block is ignored. + * + *

Currently supported type: {@code scalar_quantization_8_bit}. + */ + QuantizationConfig quantizationConfig() default @QuantizationConfig; + + /** + * Nested annotation representing {@code quantization_config} block. + */ + @interface QuantizationConfig { + /** + * Quantization type. + * + *

Currently supported: {@code scalar_quantization_8_bit}. + */ + String quantizationType() default ""; + + /** + * Quantile for scalar quantization. + */ + float quantile() default -1.0f; + + /** + * Sample size for estimating quantile(s). + */ + int sampleSize() default 20_000; + + /** + * Minimal number of points in the index required to enable quantization. + */ + int quantizationThreshold() default 100_000; + } + } diff --git a/src/main/java/ru/rt/restream/reindexer/vector/HnswConfig.java b/src/main/java/ru/rt/restream/reindexer/vector/HnswConfig.java index 591fce8..9de53dd 100644 --- a/src/main/java/ru/rt/restream/reindexer/vector/HnswConfig.java +++ b/src/main/java/ru/rt/restream/reindexer/vector/HnswConfig.java @@ -63,4 +63,47 @@ public class HnswConfig implements IndexConfig { * Enable multithreading insert mode. */ private int multithreading; + + /** + * Optional quantization configuration for HNSW index. + * + *

+ * When this value is {@code null} the {@code quantization_config} block is not + * serialized. + */ + private QuantizationConfig quantizationConfig; + + @Setter + @Getter + @EqualsAndHashCode + @NoArgsConstructor(access = AccessLevel.PACKAGE) + public static class QuantizationConfig { + /** + * Quantization type. + * + *

+ * Currently supported: {@code scalar_quantization_8_bit}. + */ + private String quantizationType; + + /** + * Quantile for scalar quantization. + * + *

+ * If this field is {@code null}, quantile is expected to be computed + * automatically + * by Reindexer. + */ + private Float quantile; + + /** + * Sample size for estimating quantile(s). + */ + private int sampleSize; + + /** + * Minimal number of samples/points required to enable quantization. + */ + private int quantizationThreshold; + } } diff --git a/src/main/java/ru/rt/restream/reindexer/vector/HnswConfigs.java b/src/main/java/ru/rt/restream/reindexer/vector/HnswConfigs.java index 1c8a10d..f0cdaf7 100644 --- a/src/main/java/ru/rt/restream/reindexer/vector/HnswConfigs.java +++ b/src/main/java/ru/rt/restream/reindexer/vector/HnswConfigs.java @@ -30,6 +30,11 @@ public class HnswConfigs { new IllegalArgumentException("HNSW index should have 'm' parameter in range [2, 128]"); public static final IllegalArgumentException EF_CONSTRUCTION_NOT_IN_RANGE_EX = new IllegalArgumentException("HNSW index should have 'efConstruction' parameter in range [4, 1024]"); + public static final IllegalArgumentException QUANTIZATION_TYPE_NOT_SUPPORTED_EX = + new IllegalArgumentException("Only 'scalar_quantization_8_bit' is supported for HNSW quantization_config."); + public static final IllegalArgumentException QUANTILE_OUT_OF_RANGE_EX = + new IllegalArgumentException("Quantile for scalar quantization must be in range [0.95, 1.0]."); + private static final float QUANTILE_UNSET = -1.0f; public static HnswConfig of(Hnsw annotation) { if (annotation.metric() == null) { @@ -52,6 +57,28 @@ public static HnswConfig of(Hnsw annotation) { config.setM(annotation.m()); config.setEfConstruction(annotation.efConstruction()); config.setMultithreading(annotation.multithreading() ? 1 : 0); + + String quantizationType = annotation.quantizationConfig().quantizationType(); + float quantile = annotation.quantizationConfig().quantile(); + int sampleSize = annotation.quantizationConfig().sampleSize(); + int quantizationThreshold = annotation.quantizationConfig().quantizationThreshold(); + + if (quantizationType != null && !quantizationType.isEmpty()) { + if (!"scalar_quantization_8_bit".equals(quantizationType)) { + throw QUANTIZATION_TYPE_NOT_SUPPORTED_EX; + } + HnswConfig.QuantizationConfig quantizationConfig = new HnswConfig.QuantizationConfig(); + quantizationConfig.setQuantizationType(quantizationType); + if (quantile != QUANTILE_UNSET) { + if (quantile < 0.95f || quantile > 1.0f) { + throw QUANTILE_OUT_OF_RANGE_EX; + } + quantizationConfig.setQuantile(quantile); + } + quantizationConfig.setSampleSize(sampleSize); + quantizationConfig.setQuantizationThreshold(quantizationThreshold); + config.setQuantizationConfig(quantizationConfig); + } return config; } diff --git a/src/test/java/ru/rt/restream/reindexer/connector/BuiltinHnswQuantizationConfigUpdateTest.java b/src/test/java/ru/rt/restream/reindexer/connector/BuiltinHnswQuantizationConfigUpdateTest.java new file mode 100644 index 0000000..a6e4031 --- /dev/null +++ b/src/test/java/ru/rt/restream/reindexer/connector/BuiltinHnswQuantizationConfigUpdateTest.java @@ -0,0 +1,23 @@ +/* + * Copyright 2020 Restream + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package ru.rt.restream.reindexer.connector; + +import ru.rt.restream.category.BuiltinTest; + +@BuiltinTest +public class BuiltinHnswQuantizationConfigUpdateTest extends HnswQuantizationConfigUpdateTest { +} + diff --git a/src/test/java/ru/rt/restream/reindexer/connector/CprotoHnswQuantizationConfigUpdateTest.java b/src/test/java/ru/rt/restream/reindexer/connector/CprotoHnswQuantizationConfigUpdateTest.java new file mode 100644 index 0000000..eeb07db --- /dev/null +++ b/src/test/java/ru/rt/restream/reindexer/connector/CprotoHnswQuantizationConfigUpdateTest.java @@ -0,0 +1,23 @@ +/* + * Copyright 2020 Restream + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package ru.rt.restream.reindexer.connector; + +import ru.rt.restream.category.CprotoTest; + +@CprotoTest +public class CprotoHnswQuantizationConfigUpdateTest extends HnswQuantizationConfigUpdateTest { +} + diff --git a/src/test/java/ru/rt/restream/reindexer/connector/HnswQuantizationConfigUpdateTest.java b/src/test/java/ru/rt/restream/reindexer/connector/HnswQuantizationConfigUpdateTest.java new file mode 100644 index 0000000..def5679 --- /dev/null +++ b/src/test/java/ru/rt/restream/reindexer/connector/HnswQuantizationConfigUpdateTest.java @@ -0,0 +1,198 @@ +/* + * Copyright 2020 Restream + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package ru.rt.restream.reindexer.connector; + +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.Setter; +import org.junit.jupiter.api.Test; +import ru.rt.restream.reindexer.CollateMode; +import ru.rt.restream.reindexer.FieldType; +import ru.rt.restream.reindexer.IndexType; +import ru.rt.restream.reindexer.Namespace; +import ru.rt.restream.reindexer.NamespaceOptions; +import ru.rt.restream.reindexer.ReindexerIndex; +import ru.rt.restream.reindexer.annotations.Hnsw; +import ru.rt.restream.reindexer.annotations.Json; +import ru.rt.restream.reindexer.annotations.Metric; +import ru.rt.restream.reindexer.annotations.Reindex; +import ru.rt.restream.reindexer.db.DbBaseTest; +import ru.rt.restream.reindexer.vector.HnswConfig; + +import java.lang.reflect.Constructor; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.nullValue; +import static org.junit.jupiter.api.Assertions.assertEquals; + +abstract class HnswQuantizationConfigUpdateTest extends DbBaseTest { + private final String namespaceName = "hnsw_quantization_config_update"; + private final String vectorIndexName = "vector"; + + @Test + void testChangeQuantizationConfig() { + db.openNamespace(namespaceName, NamespaceOptions.defaultOptions(), VectorItem.class); + + NamespaceDescriptionResponse before = getNamespaceDescription(namespaceName); + IndexResponse vectorIndexBefore = getIndexByName(before.getIndexes(), vectorIndexName); + QuantizationConfigResponse quantBefore = vectorIndexBefore.getConfig().getQuantizationConfig(); + + assertThat(quantBefore.getQuantizationType(), is("scalar_quantization_8_bit")); + assertThat(quantBefore.getSampleSize(), is(20_000)); + assertThat(quantBefore.getQuantizationThreshold(), is(100_000)); + assertThat(quantBefore.getQuantile(), nullValue()); + + ReindexerIndex updIndex = createHnswIndexForUpdate(vectorIndexName); + HnswConfig hnswConfig = createHnswConfigForUpdate(); + updIndex.setConfig(hnswConfig); + db.updateIndex(namespaceName, updIndex); + + NamespaceDescriptionResponse after = getNamespaceDescription(namespaceName); + IndexResponse vectorIndexAfter = getIndexByName(after.getIndexes(), vectorIndexName); + QuantizationConfigResponse quantAfter = vectorIndexAfter.getConfig().getQuantizationConfig(); + + assertThat(quantAfter.getQuantizationType(), is("scalar_quantization_8_bit")); + assertEquals(0.987f, quantAfter.getQuantile(), 1e-6f); + assertThat(quantAfter.getSampleSize(), is(12340)); + assertThat(quantAfter.getQuantizationThreshold(), is(43210)); + } + + private NamespaceDescriptionResponse getNamespaceDescription(String nsName) { + Namespace serviceNamespace = db.openNamespace("#namespaces", + NamespaceOptions.defaultOptions(), NamespaceDescriptionResponse.class); + String query = String.format("select * from #namespaces where name = '%s'", nsName); + Iterator iterator = serviceNamespace.execSql(query); + if (!iterator.hasNext()) { + throw new AssertionError("Namespace is not found: " + nsName); + } + return iterator.next(); + } + + private IndexResponse getIndexByName(List indexes, String indexName) { + return indexes.stream() + .filter(i -> indexName.equals(i.getName())) + .findFirst() + .orElseThrow(() -> new AssertionError("Index is not found: " + indexName)); + } + + private ReindexerIndex createHnswIndexForUpdate(String indexName) { + // index update requires the same structural info (index type + field type + + // json path) + return ReindexerIndex.builder() + .name(indexName) + .jsonPaths(Collections.singletonList(indexName)) + .indexType(IndexType.HNSW) + .fieldType(FieldType.FLOAT_VECTOR) + .collateMode(CollateMode.NONE) + .build(); + } + + private HnswConfig createHnswConfigForUpdate() { + HnswConfig config = newHnswConfig(); + config.setMetric("cosine"); + config.setDimension(1024); + config.setStartSize(1000); + config.setM(20); + config.setEfConstruction(150); + config.setMultithreading(0); + + HnswConfig.QuantizationConfig quantConfig = newQuantizationConfig(); + quantConfig.setQuantizationType("scalar_quantization_8_bit"); + quantConfig.setQuantile(0.987f); + quantConfig.setSampleSize(12340); + quantConfig.setQuantizationThreshold(43210); + config.setQuantizationConfig(quantConfig); + return config; + } + + private static HnswConfig newHnswConfig() { + try { + Constructor ctor = HnswConfig.class.getDeclaredConstructor(); + ctor.setAccessible(true); + return ctor.newInstance(); + } catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + } + + private static HnswConfig.QuantizationConfig newQuantizationConfig() { + try { + Constructor ctor = HnswConfig.QuantizationConfig.class + .getDeclaredConstructor(); + ctor.setAccessible(true); + return ctor.newInstance(); + } catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + } + + @Getter + @Setter + public static class NamespaceDescriptionResponse { + private String name; + private List indexes; + } + + @Getter + @Setter + public static class IndexResponse { + private String name; + + @Json("config") + private IndexConfigResponse config; + } + + @Getter + @Setter + public static class IndexConfigResponse { + @Json("quantization_config") + private QuantizationConfigResponse quantizationConfig; + } + + @Getter + @Setter + public static class QuantizationConfigResponse { + @Json("quantization_type") + private String quantizationType; + + @Json("quantile") + private Float quantile; + + @Json("sample_size") + private Integer sampleSize; + + @Json("quantization_threshold") + private Integer quantizationThreshold; + } + + @Getter + @Setter + @NoArgsConstructor + @AllArgsConstructor + public static class VectorItem { + @Reindex(name = "id", isPrimaryKey = true) + private Integer id; + + @Reindex(name = "vector") + @Hnsw(metric = Metric.COSINE, dimension = 1024, quantizationConfig = @Hnsw.QuantizationConfig(quantizationType = "scalar_quantization_8_bit")) + private float[] vector; + } +} diff --git a/src/test/java/ru/rt/restream/reindexer/fast/ReindexAnnotationScannerHnswIndexTest.java b/src/test/java/ru/rt/restream/reindexer/fast/ReindexAnnotationScannerHnswIndexTest.java index 3e81396..d34e714 100644 --- a/src/test/java/ru/rt/restream/reindexer/fast/ReindexAnnotationScannerHnswIndexTest.java +++ b/src/test/java/ru/rt/restream/reindexer/fast/ReindexAnnotationScannerHnswIndexTest.java @@ -62,6 +62,7 @@ public void testConfigDefaultValues_isOk() { assertThat(vector.getHnswConfig().getM(), is(16)); assertThat(vector.getHnswConfig().getEfConstruction(), is(200)); assertThat(vector.getHnswConfig().getMultithreading(), is(0)); + assertThat(vector.getHnswConfig().getQuantizationConfig(), nullValue()); } @Test @@ -84,6 +85,7 @@ public void testConfigExplicitlySetValues_isOk() { assertThat(vector.getHnswConfig().getM(), is(5)); assertThat(vector.getHnswConfig().getEfConstruction(), is(4)); assertThat(vector.getHnswConfig().getMultithreading(), is(1)); + assertThat(vector.getHnswConfig().getQuantizationConfig(), nullValue()); } @Test @@ -154,6 +156,21 @@ static class ItemWithAllConfigValues { private float[] vector; } + @Test + public void testConfigScalarQuantization8Bit_isOk() { + List indexes = scanner.parseIndexes(ItemWithScalarQuantization8BitConfig.class); + ReindexerIndex vector = getIndexByName(indexes, "hnsw_vector"); + + assertThat(vector.getHnswConfig(), notNullValue()); + assertThat(vector.getHnswConfig().getQuantizationConfig(), notNullValue()); + + assertThat(vector.getHnswConfig().getQuantizationConfig().getQuantizationType(), + is("scalar_quantization_8_bit")); + assertThat(vector.getHnswConfig().getQuantizationConfig().getQuantile(), is(0.987f)); + assertThat(vector.getHnswConfig().getQuantizationConfig().getSampleSize(), is(3000)); + assertThat(vector.getHnswConfig().getQuantizationConfig().getQuantizationThreshold(), is(5000)); + } + static class ItemWithNotHnswIndex { @Reindex(name = "id", isPrimaryKey = true) private Integer id; @@ -199,4 +216,12 @@ static class ItemWithoutHnswAnnotation { private float[] vector; } + static class ItemWithScalarQuantization8BitConfig { + @Reindex(name = "id", isPrimaryKey = true) + private Integer id; + + @Reindex(name = "hnsw_vector", type = HNSW) + @Hnsw(metric = Metric.COSINE, dimension = 1024, quantizationConfig = @Hnsw.QuantizationConfig(quantizationType = "scalar_quantization_8_bit", quantile = 0.987f, sampleSize = 3000, quantizationThreshold = 5000)) + private float[] vector; + } }