|
| 1 | +/* |
| 2 | + * Copyright OpenSearch Contributors |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + */ |
| 5 | + |
| 6 | +package org.opensearch.knn.index.mapper; |
| 7 | + |
| 8 | +import lombok.AllArgsConstructor; |
| 9 | +import lombok.Getter; |
| 10 | +import org.opensearch.core.common.Strings; |
| 11 | + |
| 12 | +import java.util.Arrays; |
| 13 | +import java.util.Locale; |
| 14 | +import java.util.stream.Collectors; |
| 15 | + |
| 16 | +/** |
| 17 | + * Enum representing the compression level for float vectors. Compression in this sense refers to compressing a |
| 18 | + * full precision value into a smaller number of bits. For instance. "16x" compression would mean that 2 bits would |
| 19 | + * need to be used to represent a 32-bit floating point number. |
| 20 | + */ |
| 21 | +@AllArgsConstructor |
| 22 | +public enum CompressionLevel { |
| 23 | + NOT_CONFIGURED(-1, ""), |
| 24 | + x1(1, "1x"), |
| 25 | + x2(2, "2x"), |
| 26 | + x4(4, "4x"), |
| 27 | + x8(8, "8x"), |
| 28 | + x16(16, "16x"), |
| 29 | + x32(32, "32x"); |
| 30 | + |
| 31 | + // Internally, an empty string is easier to deal with them null. However, from the mapping, |
| 32 | + // we do not want users to pass in the empty string and instead want null. So we make the conversion herex |
| 33 | + static final String[] NAMES_ARRAY = Arrays.stream(CompressionLevel.values()) |
| 34 | + .map(compressionLevel -> compressionLevel == NOT_CONFIGURED ? null : compressionLevel.getName()) |
| 35 | + .collect(Collectors.toList()) |
| 36 | + .toArray(new String[0]); |
| 37 | + |
| 38 | + /** |
| 39 | + * Default is set to 1x and is a noop |
| 40 | + */ |
| 41 | + private static final CompressionLevel DEFAULT = x1; |
| 42 | + |
| 43 | + /** |
| 44 | + * Get the compression level from a string representation. The format for the string should be "Nx", where N is |
| 45 | + * the factor by which compression should take place |
| 46 | + * |
| 47 | + * @param name String representation of the compression level |
| 48 | + * @return CompressionLevel enum value |
| 49 | + */ |
| 50 | + public static CompressionLevel fromName(String name) { |
| 51 | + if (Strings.isEmpty(name)) { |
| 52 | + return NOT_CONFIGURED; |
| 53 | + } |
| 54 | + for (CompressionLevel config : CompressionLevel.values()) { |
| 55 | + if (config.getName() != null && config.getName().equals(name)) { |
| 56 | + return config; |
| 57 | + } |
| 58 | + } |
| 59 | + throw new IllegalArgumentException(String.format(Locale.ROOT, "Invalid compression level: \"[%s]\"", name)); |
| 60 | + } |
| 61 | + |
| 62 | + private final int compressionLevel; |
| 63 | + @Getter |
| 64 | + private final String name; |
| 65 | + |
| 66 | + /** |
| 67 | + * Gets the number of bits used to represent a float in order to achieve this compression. For instance, for |
| 68 | + * 32x compression, each float would need to be encoded in a single bit. |
| 69 | + * |
| 70 | + * @return number of bits to represent a float at this compression level |
| 71 | + */ |
| 72 | + public int numBitsForFloat32() { |
| 73 | + if (this == NOT_CONFIGURED) { |
| 74 | + return DEFAULT.numBitsForFloat32(); |
| 75 | + } |
| 76 | + |
| 77 | + return (Float.BYTES * Byte.SIZE) / compressionLevel; |
| 78 | + } |
| 79 | + |
| 80 | + /** |
| 81 | + * Utility method that checks if compression is configured. |
| 82 | + * |
| 83 | + * @param compressionLevel Compression to check |
| 84 | + * @return true if compression is configured, false otherwise |
| 85 | + */ |
| 86 | + public static boolean isConfigured(CompressionLevel compressionLevel) { |
| 87 | + return compressionLevel != null && compressionLevel != NOT_CONFIGURED; |
| 88 | + } |
| 89 | +} |
0 commit comments