This is an automated email from the ASF dual-hosted git repository. asf-gitbox-commits pushed a commit to branch PHOENIX-7998 in repository https://gitbox.apache.org/repos/asf/phoenix.git
commit 9f99371a58d17b53fe695e51a10a46e716644b84 Author: Andrew Purtell <[email protected]> AuthorDate: Fri Sep 18 23:25:41 2026 -0700 PHOENIX-7998 Phase 6: IVF centroid management Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]> --- .../apache/phoenix/cache/VectorCentroidCache.java | 319 +++++++++++ .../phoenix/index/vector/CentroidManager.java | 225 ++++++++ .../phoenix/index/vector/ClusterSkewMetrics.java | 250 +++++++++ .../apache/phoenix/index/vector/KMeansConfig.java | 147 ++++++ .../apache/phoenix/index/vector/KMeansResult.java | 131 +++++ .../apache/phoenix/index/vector/KMeansTrainer.java | 358 +++++++++++++ .../phoenix/index/vector/VectorIndexTrainer.java | 213 ++++++++ .../apache/phoenix/optimize/QueryOptimizer.java | 7 +- .../org/apache/phoenix/query/QueryServices.java | 3 + .../apache/phoenix/query/QueryServicesOptions.java | 1 + .../org/apache/phoenix/schema/MetaDataClient.java | 45 +- .../java/org/apache/phoenix/util/IndexUtil.java | 2 +- .../phoenix/coprocessor/MetaDataEndpointImpl.java | 18 +- .../phoenix/end2end/VectorCentroidTableIT.java | 100 ++++ .../org/apache/phoenix/end2end/VectorIndexIT.java | 157 ++++++ .../phoenix/cache/VectorCentroidCacheTest.java | 315 +++++++++++ .../phoenix/index/vector/CentroidManagerTest.java | 65 +++ .../phoenix/index/vector/KMeansTrainerTest.java | 585 +++++++++++++++++++++ .../index/vector/VectorIndexTrainerTest.java | 109 ++++ .../phoenix/optimize/VectorSearchUtilTest.java | 26 + 20 files changed, 3051 insertions(+), 25 deletions(-) diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/cache/VectorCentroidCache.java b/phoenix-core-client/src/main/java/org/apache/phoenix/cache/VectorCentroidCache.java new file mode 100644 index 0000000000..0e643c5955 --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/cache/VectorCentroidCache.java @@ -0,0 +1,319 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.cache; + +import java.sql.Connection; +import java.sql.SQLException; +import java.util.List; +import java.util.Objects; +import java.util.PriorityQueue; +import java.util.concurrent.ExecutionException; +import org.apache.hadoop.conf.Configuration; +import org.apache.phoenix.exception.SQLExceptionCode; +import org.apache.phoenix.exception.SQLExceptionInfo; +import org.apache.phoenix.index.vector.CentroidManager; +import org.apache.phoenix.index.vector.KMeansTrainer; +import org.apache.phoenix.optimize.DistanceMetric; +import org.apache.phoenix.query.QueryServices; +import org.apache.phoenix.query.QueryServicesOptions; +import org.apache.phoenix.schema.types.PDataType; +import org.apache.phoenix.schema.types.PVectorDouble; +import org.apache.phoenix.schema.types.PVectorFloat; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.apache.phoenix.thirdparty.com.google.common.cache.Cache; +import org.apache.phoenix.thirdparty.com.google.common.cache.CacheBuilder; +import org.apache.phoenix.thirdparty.com.google.common.cache.CacheStats; + +/** + * Process-wide LRU cache of IVF centroid models, keyed by index name and generation ID. The model + * of a generation does not change, so the cache does not refresh an entry. A byte limit on the + * centroid vectors controls the size of the cache. A miss loads the model from + * {@code SYSTEM.VECTOR_CENTROID}, and concurrent misses for the same key share one load. + */ +public final class VectorCentroidCache { + + private static final Logger LOGGER = LoggerFactory.getLogger(VectorCentroidCache.class); + + private static volatile VectorCentroidCache instance; + + private final Cache<Key, CachedCentroids> cache; + private final long maxBytes; + + VectorCentroidCache(long maxBytes) { + this.maxBytes = maxBytes; + // Use one segment, so that the byte limit applies to the whole cache. Guava divides a weight + // limit equally among segments, and at insert it evicts a model heavier than one share. + this.cache = CacheBuilder.newBuilder().maximumWeight(maxBytes) + .weigher((Key key, CachedCentroids centroids) -> centroids.getWeight()).concurrencyLevel(1) + .recordStats().build(); + } + + /** + * Returns the process-wide cache. The first call sets the byte limit from {@code conf}, and later + * calls ignore {@code conf}. + */ + public static VectorCentroidCache getInstance(Configuration conf) { + VectorCentroidCache result = instance; + if (result == null) { + synchronized (VectorCentroidCache.class) { + result = instance; + if (result == null) { + instance = result = new VectorCentroidCache( + conf.getLong(QueryServices.VECTOR_CENTROID_CACHE_MAX_BYTES_ATTRIB, + QueryServicesOptions.DEFAULT_VECTOR_CENTROID_CACHE_MAX_BYTES)); + } + } + } + return result; + } + + /** + * Returns the centroids of the index generation, and loads them on a cache miss. + * @throws SQLException if no centroids are recorded for the generation, or if the read fails. The + * cache does not keep a failed load, so a later call tries again. + */ + public CachedCentroids get(Connection conn, String indexName, long generation, + DistanceMetric metric) throws SQLException { + Key key = new Key(indexName, generation); + try { + return cache.get(key, () -> checkWeight(key, load(conn, key, metric))); + } catch (ExecutionException e) { + if (e.getCause() instanceof SQLException) { + throw (SQLException) e.getCause(); + } + throw new SQLException(e.getCause()); + } + } + + /** Returns the cached centroids of the index generation, or null on a miss. It does not load. */ + public CachedCentroids getIfPresent(String indexName, long generation) { + return cache.getIfPresent(new Key(indexName, generation)); + } + + /** Puts a trained model into the cache, so that the process does not load it again. */ + public void put(String indexName, long generation, CachedCentroids centroids) { + Key key = new Key(indexName, generation); + cache.put(key, checkWeight(key, centroids)); + } + + /** Removes all cached generations of the index. */ + public void invalidate(String indexName) { + cache.asMap().keySet().removeIf(k -> k.indexName.equals(indexName)); + } + + public long size() { + return cache.size(); + } + + public CacheStats stats() { + return cache.stats(); + } + + /** + * Logs a warning if the model is heavier than the whole cache. The cache evicts such a model at + * insert, so each use loads it again. + */ + private CachedCentroids checkWeight(Key key, CachedCentroids centroids) { + if (centroids.getWeight() > maxBytes) { + LOGGER.warn( + "Centroids of vector index {} generation {} take {} bytes, more than {} = {}, so the" + + " cache cannot retain them and queries reload them on each use", + key.indexName, key.generation, centroids.getWeight(), + QueryServices.VECTOR_CENTROID_CACHE_MAX_BYTES_ATTRIB, maxBytes); + } + return centroids; + } + + private static CachedCentroids load(Connection conn, Key key, DistanceMetric metric) + throws SQLException { + List<float[]> centroids = CentroidManager.loadCentroids(conn, key.indexName, key.generation); + if (centroids.isEmpty()) { + throw new SQLExceptionInfo.Builder(SQLExceptionCode.INVALID_VECTOR_INDEX_PARAMS).setMessage( + "No centroids recorded for vector index " + key.indexName + " generation " + key.generation) + .build().buildException(); + } + return new CachedCentroids(centroids, metric); + } + + private static final class Key { + private final String indexName; + private final long generation; + + Key(String indexName, long generation) { + this.indexName = Objects.requireNonNull(indexName, "indexName"); + this.generation = generation; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Key)) { + return false; + } + Key other = (Key) o; + return generation == other.generation && indexName.equals(other.indexName); + } + + @Override + public int hashCode() { + return 31 * indexName.hashCode() + Long.hashCode(generation); + } + } + + /** + * Immutable centroid model of one generation. It assigns a vector to its nearest centroid, and it + * selects the nearest centroids that a query probes. It uses + * {@link KMeansTrainer#assignmentDistance}, so training, writes and probes route a vector the + * same way. + */ + public static final class CachedCentroids { + private final float[][] centroids; + private final int dimension; + private final DistanceMetric metric; + + public CachedCentroids(List<float[]> centroids, DistanceMetric metric) { + if (centroids.isEmpty()) { + throw new IllegalArgumentException("centroids must not be empty"); + } + this.centroids = centroids.toArray(new float[0][]); + this.dimension = this.centroids[0].length; + this.metric = Objects.requireNonNull(metric, "metric"); + } + + public int size() { + return centroids.length; + } + + public int getDimension() { + return dimension; + } + + /** Returns the size of the centroid vectors in bytes, at most Integer.MAX_VALUE. */ + int getWeight() { + return (int) Math.min(Integer.MAX_VALUE, (long) centroids.length * dimension * Float.BYTES); + } + + public DistanceMetric getMetric() { + return metric; + } + + /** Returns the centroid vector that has the specified centroid ID. */ + public float[] getCentroid(int id) { + return centroids[id]; + } + + /** + * Returns the ID of the centroid nearest to the vector. Equal distances resolve to the lowest + * centroid ID. + */ + public int assign(float[] vector) { + checkDimension(vector.length); + vector = saturate(vector); + int best = 0; + double bestDist = Double.MAX_VALUE; + for (int c = 0; c < centroids.length; c++) { + double d = KMeansTrainer.assignmentDistance(metric, vector, centroids[c]); + if (d < bestDist) { + bestDist = d; + best = c; + } + } + return best; + } + + /** Returns the ID of the centroid nearest to a packed FLOAT or DOUBLE vector value. */ + public int assign(byte[] buf, int offset, int length, PDataType<?> vectorType) { + return assign(decode(buf, offset, length, vectorType)); + } + + /** + * Returns the IDs of the {@code n} centroids nearest to {@code vector}, in ascending distance + * order. Equal distances sort by ascending ID, so the first ID is the ID that {@link #assign} + * returns. + */ + public int[] nearest(float[] vector, int n) { + checkDimension(vector.length); + vector = saturate(vector); + int count = Math.min(n, centroids.length); + double[] dist = new double[centroids.length]; + PriorityQueue<Integer> farthestFirst = new PriorityQueue<>(count + 1, (a, b) -> { + int cmp = Double.compare(dist[b], dist[a]); + return cmp != 0 ? cmp : Integer.compare(b, a); + }); + for (int c = 0; c < centroids.length; c++) { + dist[c] = KMeansTrainer.assignmentDistance(metric, vector, centroids[c]); + farthestFirst.add(c); + if (farthestFirst.size() > count) { + farthestFirst.poll(); + } + } + int[] result = new int[count]; + for (int i = count - 1; i >= 0; i--) { + result[i] = farthestFirst.poll(); + } + return result; + } + + /** + * Decodes packed vector bytes into a float array. A DOUBLE element beyond float range becomes + * infinite. + */ + public static float[] decode(byte[] buf, int offset, int length, PDataType<?> vectorType) { + if (vectorType == PVectorDouble.INSTANCE) { + double[] d = PVectorDouble.readElements(buf, offset, length); + float[] f = new float[d.length]; + for (int i = 0; i < d.length; i++) { + f[i] = (float) d[i]; + } + return f; + } + return PVectorFloat.readElements(buf, offset, length); + } + + /** + * Replaces each infinite component with the largest finite float of the same sign. A finite + * DOUBLE element beyond float range narrows to an infinite component. Under COSINE, the vector + * then routes by its direction. Under L2 and INNER_PRODUCT, the saturated component dominates + * every distance, so the distances tie. {@link #assign} and {@link #nearest} resolve the tie to + * the lowest centroid ID, so writes, probes and inspection route the vector the same way. + */ + private static float[] saturate(float[] vector) { + float[] out = vector; + for (int i = 0; i < vector.length; i++) { + if (Float.isInfinite(vector[i])) { + if (out == vector) { + out = vector.clone(); + } + out[i] = Math.copySign(Float.MAX_VALUE, vector[i]); + } + } + return out; + } + + private void checkDimension(int length) { + if (length != dimension) { + throw new IllegalArgumentException( + "Vector dimension " + length + " does not match centroid dimension " + dimension); + } + } + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/CentroidManager.java b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/CentroidManager.java new file mode 100644 index 0000000000..d76f0463e8 --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/CentroidManager.java @@ -0,0 +1,225 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.index.vector; + +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.CENTROID_ID; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.CENTROID_VECTOR; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.GENERATION_ID; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.INDEX_NAME; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.INDEX_STATE; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.SYSTEM_CATALOG_SCHEMA; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.SYSTEM_CATALOG_TABLE; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.SYSTEM_VECTOR_CENTROID_NAME; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.TABLE_NAME; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.TABLE_SCHEM; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.TENANT_ID; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.VECTOR_CENTROID_GENERATION; +import static org.apache.phoenix.jdbc.PhoenixDatabaseMetaData.VECTOR_IVF_LISTS; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.List; +import java.util.Properties; +import org.apache.hadoop.hbase.client.Mutation; +import org.apache.phoenix.coprocessorclient.MetaDataProtocol.MetaDataMutationResult; +import org.apache.phoenix.coprocessorclient.MetaDataProtocol.MutationCode; +import org.apache.phoenix.exception.SQLExceptionCode; +import org.apache.phoenix.exception.SQLExceptionInfo; +import org.apache.phoenix.jdbc.PhoenixConnection; +import org.apache.phoenix.schema.PTable; +import org.apache.phoenix.schema.TableNotFoundException; +import org.apache.phoenix.schema.types.PVectorFloat; +import org.apache.phoenix.util.EnvironmentEdgeManager; +import org.apache.phoenix.util.PhoenixRuntime; +import org.apache.phoenix.util.SchemaUtil; + +/** + * Stores IVF centroid generations in {@code SYSTEM.VECTOR_CENTROID}, and records the generations of + * each index in {@code SYSTEM.CATALOG}. + * <p> + * Centroid rows are keyed by the full index name, which is case sensitive. Generation IDs of an + * index strictly increase. A re-created index takes its first generation ID from the wall clock. If + * the clock does not go back, its IDs are higher than the IDs of the dropped index. Thus no cache + * can serve the model of an earlier index with the same name. + * <p> + * Some methods commit or roll back the connection that they receive. A caller with uncommitted + * mutations must give them an internal connection from {@link #newInternalConnection}. + */ +public final class CentroidManager { + + private CentroidManager() { + } + + /** + * Opens an internal connection with the client properties of {@code conn}, but without the tenant + * ID and the SCN. Writes to system tables on this connection stay separate from the uncommitted + * mutations of {@code conn}. + */ + public static PhoenixConnection newInternalConnection(PhoenixConnection conn) + throws SQLException { + Properties props = new Properties(); + Properties clientInfo = conn.getClientInfo(); + for (String key : clientInfo.stringPropertyNames()) { + if ( + !PhoenixRuntime.TENANT_ID_ATTRIB.equals(key) + && !PhoenixRuntime.CURRENT_SCN_ATTRIB.equals(key) + ) { + props.setProperty(key, clientInfo.getProperty(key)); + } + } + return new PhoenixConnection(conn, conn.getQueryServices(), props); + } + + /** + * Returns a new generation ID from the wall clock. If {@code current} is not null, the ID is + * always greater than {@code current}. Thus the IDs of an index strictly increase, even if the + * clock goes back. If {@code current} is null, the ID is the wall clock time. + */ + public static long nextGeneration(Long current) { + long now = EnvironmentEdgeManager.currentTimeMillis(); + return current == null ? now : Math.max(current + 1, now); + } + + /** Writes the centroid vectors of a generation with IDs from 0, and commits the connection. */ + public static void persistCentroids(Connection conn, String indexName, long generation, + List<float[]> centroids) throws SQLException { + String sql = "UPSERT INTO " + SYSTEM_VECTOR_CENTROID_NAME + " (" + INDEX_NAME + ", " + + GENERATION_ID + ", " + CENTROID_ID + ", " + CENTROID_VECTOR + ") VALUES (?, ?, ?, ?)"; + try (PreparedStatement ps = conn.prepareStatement(sql)) { + for (int i = 0; i < centroids.size(); i++) { + ps.setString(1, indexName); + ps.setLong(2, generation); + ps.setInt(3, i); + ps.setBytes(4, PVectorFloat.INSTANCE.toBytes(centroids.get(i))); + ps.executeUpdate(); + } + } + conn.commit(); + } + + /** + * Loads the centroid vectors of a generation in ascending centroid ID order. Returns an empty + * list if the generation has no recorded centroids. + */ + public static List<float[]> loadCentroids(Connection conn, String indexName, long generation) + throws SQLException { + String sql = "SELECT " + CENTROID_VECTOR + " FROM " + SYSTEM_VECTOR_CENTROID_NAME + " WHERE " + + INDEX_NAME + " = ? AND " + GENERATION_ID + " = ? AND " + CENTROID_ID + " >= 0 AND " + + CENTROID_VECTOR + " IS NOT NULL ORDER BY " + CENTROID_ID; + List<float[]> centroids = new ArrayList<>(); + try (PreparedStatement ps = conn.prepareStatement(sql)) { + ps.setString(1, indexName); + ps.setLong(2, generation); + try (ResultSet rs = ps.executeQuery()) { + while (rs.next()) { + byte[] b = rs.getBytes(1); + centroids.add(PVectorFloat.readElements(b, 0, b.length)); + } + } + } + return centroids; + } + + /** + * Records the active centroid generation and the list count in {@code SYSTEM.CATALOG} through the + * metadata endpoint, and returns the updated index. The endpoint keeps the current index state. + * It advances the DDL timestamp, so that clients refresh their cached index metadata. This method + * rolls back the uncommitted mutations of {@code conn}. + */ + public static PTable setGenerationAndLists(PhoenixConnection conn, PTable index, long generation, + int lists) throws SQLException { + String schemaName = index.getSchemaName().getString(); + String tableName = index.getTableName().getString(); + String sql = "UPSERT INTO " + SYSTEM_CATALOG_SCHEMA + ".\"" + SYSTEM_CATALOG_TABLE + "\"(" + + TENANT_ID + "," + TABLE_SCHEM + "," + TABLE_NAME + "," + INDEX_STATE + "," + + VECTOR_CENTROID_GENERATION + "," + VECTOR_IVF_LISTS + ") VALUES (?, ?, ?, ?, ?, ?)"; + boolean autoCommit = conn.getAutoCommit(); + List<Mutation> tableMetadata; + conn.setAutoCommit(false); + try { + try (PreparedStatement ps = conn.prepareStatement(sql)) { + ps.setString(1, index.getTenantId() == null ? null : index.getTenantId().getString()); + ps.setString(2, schemaName.isEmpty() ? null : schemaName); + ps.setString(3, tableName); + ps.setString(4, index.getIndexState().getSerializedValue()); + ps.setLong(5, generation); + ps.setInt(6, lists); + ps.execute(); + } + tableMetadata = conn.getMutationState().toMutations().next().getSecond(); + } finally { + conn.rollback(); + conn.setAutoCommit(autoCommit); + } + MetaDataMutationResult result = conn.getQueryServices().updateIndexState(tableMetadata, + index.getParentName() == null ? null : index.getParentName().getString()); + if (result.getMutationCode() == MutationCode.TABLE_NOT_FOUND) { + throw new TableNotFoundException(schemaName, tableName); + } + if (result.getMutationCode() != MutationCode.TABLE_ALREADY_EXISTS) { + throw new SQLExceptionInfo.Builder(SQLExceptionCode.INVALID_INDEX_STATE_TRANSITION) + .setMessage( + "Could not record centroid generation " + generation + ": " + result.getMutationCode()) + .setSchemaName(schemaName).setTableName(tableName).build().buildException(); + } + return result.getTable(); + } + + /** Deletes all centroid rows for a specific generation. */ + public static void deleteGeneration(Connection conn, String indexName, long generation) + throws SQLException { + try (PreparedStatement ps = conn.prepareStatement("DELETE FROM " + SYSTEM_VECTOR_CENTROID_NAME + + " WHERE " + INDEX_NAME + " = ? AND " + GENERATION_ID + " = ?")) { + ps.setString(1, indexName); + ps.setLong(2, generation); + executeAutoCommitted(conn, ps); + } + } + + /** Deletes the centroid rows of all generations of the index. */ + public static void deleteAllCentroids(Connection conn, String indexName) throws SQLException { + try (PreparedStatement ps = conn.prepareStatement( + "DELETE FROM " + SYSTEM_VECTOR_CENTROID_NAME + " WHERE " + INDEX_NAME + " = ?")) { + ps.setString(1, indexName); + executeAutoCommitted(conn, ps); + } + } + + /** + * Runs a delete with autocommit on, so that the server deletes the rows. Then it restores the + * autocommit setting of the caller. + */ + private static void executeAutoCommitted(Connection conn, PreparedStatement ps) + throws SQLException { + boolean autoCommit = conn.getAutoCommit(); + conn.setAutoCommit(true); + try { + ps.executeUpdate(); + } finally { + conn.setAutoCommit(autoCommit); + } + } + + /** Returns the escaped full name of the table, for use in SQL statements. */ + static String escapedName(PTable table) { + return SchemaUtil.getEscapedFullTableName(table.getName().getString()); + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/ClusterSkewMetrics.java b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/ClusterSkewMetrics.java new file mode 100644 index 0000000000..313dc1f9b0 --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/ClusterSkewMetrics.java @@ -0,0 +1,250 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.index.vector; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.io.IOException; +import java.util.Arrays; +import java.util.Objects; + +/** + * Size distribution of k-means clusters, and its skew. The metrics are the minimum, maximum, mean, + * 95th percentile, standard deviation and coefficient of variation of the cluster sizes. A binary + * form keeps the metrics in the generation summary row of {@code SYSTEM.VECTOR_CENTROID}. + */ +public class ClusterSkewMetrics { + + private static final int SERIALIZATION_VERSION = 1; + + private final int min; + private final int max; + private final double mean; + private final double p95; + private final double stdDev; + private final double coefficientOfVariation; + private final int[] clusterSizes; + private final int totalVectors; + private final int k; + + public ClusterSkewMetrics(int min, int max, double mean, double p95, double stdDev, + double coefficientOfVariation, int[] clusterSizes, int totalVectors, int k) { + this.min = min; + this.max = max; + this.mean = mean; + this.p95 = p95; + this.stdDev = stdDev; + this.coefficientOfVariation = coefficientOfVariation; + this.clusterSizes = clusterSizes != null ? clusterSizes.clone() : new int[0]; + this.totalVectors = totalVectors; + this.k = k; + } + + /** + * Computes the distribution metrics of the cluster sizes. The 95th percentile interpolates + * linearly between the sorted sizes. + * @param clusterSizes the count of vectors in each cluster; must not be null or empty + * @return the skew metrics + */ + public static ClusterSkewMetrics compute(int[] clusterSizes) { + if (clusterSizes == null || clusterSizes.length == 0) { + throw new IllegalArgumentException("clusterSizes must not be null or empty"); + } + + int k = clusterSizes.length; + int min = Integer.MAX_VALUE; + int max = Integer.MIN_VALUE; + long total = 0; + + for (int size : clusterSizes) { + if (size < min) { + min = size; + } + if (size > max) { + max = size; + } + total += size; + } + + int totalVectors = (int) total; + double mean = (double) total / k; + + double sumSqDiff = 0.0; + for (int size : clusterSizes) { + double diff = size - mean; + sumSqDiff += diff * diff; + } + double variance = sumSqDiff / k; + double stdDev = Math.sqrt(variance); + double cv = (mean == 0.0) ? 0.0 : (stdDev / mean); + + int[] sorted = clusterSizes.clone(); + Arrays.sort(sorted); + double p95; + if (k == 1) { + p95 = sorted[0]; + } else { + double rank = 0.95 * (k - 1); + int lower = (int) Math.floor(rank); + int upper = (int) Math.ceil(rank); + if (lower == upper) { + p95 = sorted[lower]; + } else { + p95 = sorted[lower] + (rank - lower) * (sorted[upper] - sorted[lower]); + } + } + + return new ClusterSkewMetrics(min, max, mean, p95, stdDev, cv, clusterSizes, totalVectors, k); + } + + /** + * Serializes the metrics to a versioned binary form for the generation summary row. + */ + public byte[] toBytes() { + try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); + DataOutputStream dos = new DataOutputStream(baos)) { + dos.writeInt(SERIALIZATION_VERSION); + dos.writeInt(k); + dos.writeInt(totalVectors); + dos.writeInt(min); + dos.writeInt(max); + dos.writeDouble(mean); + dos.writeDouble(p95); + dos.writeDouble(stdDev); + dos.writeDouble(coefficientOfVariation); + dos.writeInt(clusterSizes.length); + for (int size : clusterSizes) { + dos.writeInt(size); + } + dos.flush(); + return baos.toByteArray(); + } catch (IOException e) { + throw new RuntimeException("Failed to serialize ClusterSkewMetrics", e); + } + } + + /** + * Deserializes the metrics from the binary form that {@link #toBytes} writes. + * @throws IllegalArgumentException if the bytes are null or empty, or have an unsupported + * version. A truncated array causes a RuntimeException. + */ + public static ClusterSkewMetrics fromBytes(byte[] bytes) { + if (bytes == null || bytes.length == 0) { + throw new IllegalArgumentException("Byte array must not be null or empty"); + } + try (ByteArrayInputStream bais = new ByteArrayInputStream(bytes); + DataInputStream dis = new DataInputStream(bais)) { + int version = dis.readInt(); + if (version != SERIALIZATION_VERSION) { + throw new IllegalArgumentException("Unsupported ClusterSkewMetrics version: " + version); + } + int k = dis.readInt(); + int totalVectors = dis.readInt(); + int min = dis.readInt(); + int max = dis.readInt(); + double mean = dis.readDouble(); + double p95 = dis.readDouble(); + double stdDev = dis.readDouble(); + double cv = dis.readDouble(); + int sizesLen = dis.readInt(); + int[] sizes = new int[sizesLen]; + for (int i = 0; i < sizesLen; i++) { + sizes[i] = dis.readInt(); + } + return new ClusterSkewMetrics(min, max, mean, p95, stdDev, cv, sizes, totalVectors, k); + } catch (IOException e) { + throw new RuntimeException("Failed to deserialize ClusterSkewMetrics", e); + } + } + + /** + * Returns true if the coefficient of variation is more than {@link #SEVERE_SKEW_CV_THRESHOLD}. + * Training logs a warning for such a model. + */ + public boolean isSevereSkew() { + return coefficientOfVariation > 1.0; + } + + public int getMin() { + return min; + } + + public int getMax() { + return max; + } + + public double getMean() { + return mean; + } + + public double getP95() { + return p95; + } + + public double getStdDev() { + return stdDev; + } + + public double getCoefficientOfVariation() { + return coefficientOfVariation; + } + + public int[] getClusterSizes() { + return clusterSizes.clone(); + } + + public int getTotalVectors() { + return totalVectors; + } + + public int getK() { + return k; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ClusterSkewMetrics)) { + return false; + } + ClusterSkewMetrics that = (ClusterSkewMetrics) o; + return min == that.min && max == that.max && Double.compare(that.mean, mean) == 0 + && Double.compare(that.p95, p95) == 0 && Double.compare(that.stdDev, stdDev) == 0 + && Double.compare(that.coefficientOfVariation, coefficientOfVariation) == 0 + && totalVectors == that.totalVectors && k == that.k + && Arrays.equals(clusterSizes, that.clusterSizes); + } + + @Override + public int hashCode() { + int result = Objects.hash(min, max, mean, p95, stdDev, coefficientOfVariation, totalVectors, k); + result = 31 * result + Arrays.hashCode(clusterSizes); + return result; + } + + @Override + public String toString() { + return "ClusterSkewMetrics{" + "k=" + k + ", totalVectors=" + totalVectors + ", min=" + min + + ", max=" + max + ", mean=" + mean + ", p95=" + p95 + ", stdDev=" + stdDev + ", cv=" + + coefficientOfVariation + ", severeSkew=" + isSevereSkew() + '}'; + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/KMeansConfig.java b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/KMeansConfig.java new file mode 100644 index 0000000000..75110a5cc0 --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/KMeansConfig.java @@ -0,0 +1,147 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.index.vector; + +import java.util.Objects; +import org.apache.phoenix.optimize.DistanceMetric; + +public final class KMeansConfig { + + public static final int DEFAULT_MAX_ITERATIONS = 100; + public static final double DEFAULT_CONVERGENCE_THRESHOLD = 1e-4; + public static final double DEFAULT_SPLIT_THRESHOLD_MULTIPLIER = 2.0; + public static final int DEFAULT_SPLIT_REBALANCE_ITERATIONS = 5; + + private final int maxIterations; + private final double convergenceThreshold; + private final DistanceMetric distanceMetric; + private final boolean enableSplitHeuristic; + private final double splitThresholdMultiplier; + private final int splitRebalanceIterations; + private final Long randomSeed; + + private KMeansConfig(Builder builder) { + this.maxIterations = builder.maxIterations; + this.convergenceThreshold = builder.convergenceThreshold; + this.distanceMetric = builder.distanceMetric; + this.enableSplitHeuristic = builder.enableSplitHeuristic; + this.splitThresholdMultiplier = builder.splitThresholdMultiplier; + this.splitRebalanceIterations = builder.splitRebalanceIterations; + this.randomSeed = builder.randomSeed; + } + + public static Builder newBuilder() { + return new Builder(); + } + + public static KMeansConfig defaultConfig() { + return new Builder().build(); + } + + public int getMaxIterations() { + return maxIterations; + } + + public double getConvergenceThreshold() { + return convergenceThreshold; + } + + public DistanceMetric getDistanceMetric() { + return distanceMetric; + } + + public boolean isEnableSplitHeuristic() { + return enableSplitHeuristic; + } + + public double getSplitThresholdMultiplier() { + return splitThresholdMultiplier; + } + + public int getSplitRebalanceIterations() { + return splitRebalanceIterations; + } + + public Long getRandomSeed() { + return randomSeed; + } + + public static final class Builder { + private int maxIterations = DEFAULT_MAX_ITERATIONS; + private double convergenceThreshold = DEFAULT_CONVERGENCE_THRESHOLD; + private DistanceMetric distanceMetric = DistanceMetric.L2; + private boolean enableSplitHeuristic = true; + private double splitThresholdMultiplier = DEFAULT_SPLIT_THRESHOLD_MULTIPLIER; + private int splitRebalanceIterations = DEFAULT_SPLIT_REBALANCE_ITERATIONS; + private Long randomSeed; + + public Builder maxIterations(int maxIterations) { + if (maxIterations <= 0) { + throw new IllegalArgumentException("maxIterations must be > 0: " + maxIterations); + } + this.maxIterations = maxIterations; + return this; + } + + public Builder convergenceThreshold(double convergenceThreshold) { + if (convergenceThreshold <= 0.0) { + throw new IllegalArgumentException( + "convergenceThreshold must be > 0: " + convergenceThreshold); + } + this.convergenceThreshold = convergenceThreshold; + return this; + } + + public Builder distanceMetric(DistanceMetric distanceMetric) { + this.distanceMetric = Objects.requireNonNull(distanceMetric, "distanceMetric"); + return this; + } + + public Builder enableSplitHeuristic(boolean enableSplitHeuristic) { + this.enableSplitHeuristic = enableSplitHeuristic; + return this; + } + + public Builder splitThresholdMultiplier(double splitThresholdMultiplier) { + if (splitThresholdMultiplier <= 0.0) { + throw new IllegalArgumentException( + "splitThresholdMultiplier must be > 0: " + splitThresholdMultiplier); + } + this.splitThresholdMultiplier = splitThresholdMultiplier; + return this; + } + + public Builder splitRebalanceIterations(int splitRebalanceIterations) { + if (splitRebalanceIterations < 0) { + throw new IllegalArgumentException( + "splitRebalanceIterations must be >= 0: " + splitRebalanceIterations); + } + this.splitRebalanceIterations = splitRebalanceIterations; + return this; + } + + public Builder randomSeed(long randomSeed) { + this.randomSeed = randomSeed; + return this; + } + + public KMeansConfig build() { + return new KMeansConfig(this); + } + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/KMeansResult.java b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/KMeansResult.java new file mode 100644 index 0000000000..fd93e647f6 --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/KMeansResult.java @@ -0,0 +1,131 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.index.vector; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.phoenix.optimize.DistanceMetric; + +/** + * Result of k-means training: the trained centroids, the cluster of each sample vector, and the + * skew metrics before and after a cluster split. + */ +public final class KMeansResult { + + private final List<float[]> centroids; + private final int requestedK; + private final int iterations; + private final boolean converged; + private final double finalDistortion; + private final int[] clusterSizes; + private final int[] assignments; + private final ClusterSkewMetrics skewMetrics; + private final ClusterSkewMetrics preSplitSkewMetrics; + private final ClusterSkewMetrics postSplitSkewMetrics; + private final int dimension; + private final DistanceMetric distanceMetric; + private final List<Double> distortionHistory; + + KMeansResult(List<float[]> centroids, int requestedK, int iterations, boolean converged, + double finalDistortion, int[] clusterSizes, int[] assignments, ClusterSkewMetrics skewMetrics, + ClusterSkewMetrics preSplitSkewMetrics, ClusterSkewMetrics postSplitSkewMetrics, int dimension, + DistanceMetric distanceMetric, List<Double> distortionHistory) { + this.centroids = Collections.unmodifiableList(new ArrayList<>(centroids)); + this.requestedK = requestedK; + this.iterations = iterations; + this.converged = converged; + this.finalDistortion = finalDistortion; + this.clusterSizes = clusterSizes.clone(); + this.assignments = assignments.clone(); + this.skewMetrics = skewMetrics; + this.preSplitSkewMetrics = preSplitSkewMetrics; + this.postSplitSkewMetrics = postSplitSkewMetrics; + this.dimension = dimension; + this.distanceMetric = distanceMetric; + this.distortionHistory = Collections.unmodifiableList(new ArrayList<>(distortionHistory)); + } + + public List<float[]> getCentroids() { + return centroids; + } + + /** + * Returns the number of trained centroids. A cluster split can make this number larger than the + * requested count. + */ + public int getEffectiveK() { + return centroids.size(); + } + + public int getRequestedK() { + return requestedK; + } + + public int getIterations() { + return iterations; + } + + public boolean isConverged() { + return converged; + } + + public double getFinalDistortion() { + return finalDistortion; + } + + /** Returns the number of sample vectors in each cluster, by position in the centroid list. */ + public int[] getClusterSizes() { + return clusterSizes.clone(); + } + + /** Returns the position in the centroid list of the cluster of each sample vector. */ + public int[] getAssignments() { + return assignments.clone(); + } + + /** Returns the skew metrics of the final cluster sizes, after a split if one occurred. */ + public ClusterSkewMetrics getSkewMetrics() { + return skewMetrics; + } + + public ClusterSkewMetrics getPreSplitSkewMetrics() { + return preSplitSkewMetrics; + } + + /** Returns the skew metrics after a cluster split, or null if no cluster was split. */ + public ClusterSkewMetrics getPostSplitSkewMetrics() { + return postSplitSkewMetrics; + } + + public boolean hasSplit() { + return postSplitSkewMetrics != null; + } + + public int getDimension() { + return dimension; + } + + public DistanceMetric getDistanceMetric() { + return distanceMetric; + } + + public List<Double> getDistortionHistory() { + return distortionHistory; + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/KMeansTrainer.java b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/KMeansTrainer.java new file mode 100644 index 0000000000..8c93ad3490 --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/KMeansTrainer.java @@ -0,0 +1,358 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.index.vector; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Iterator; +import java.util.List; +import java.util.Random; +import org.apache.phoenix.expression.function.VectorDistanceUtil; +import org.apache.phoenix.optimize.DistanceMetric; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * In-memory k-means trainer for IVF index centroids. + * <p> + * Training starts from k-means++ seeds and runs Lloyd iterations. Under L2 and INNER_PRODUCT, each + * centroid moves to the mean of its members. Under COSINE, training uses spherical k-means and + * scales each centroid to unit length. Training reseeds empty clusters, and stops when the relative + * change in distortion is below the threshold. After training, it can split overloaded clusters to + * balance the posting lists. Training, index maintenance and query probes all use + * {@link #assignmentDistance}, so they route a vector to the same centroid. + */ +public final class KMeansTrainer { + + private static final Logger LOGGER = LoggerFactory.getLogger(KMeansTrainer.class); + + private KMeansTrainer() { + } + + /** + * Returns the distance that assigns a vector to a centroid. COSINE indexes use cosine distance. + * L2 and INNER_PRODUCT indexes use squared Euclidean distance. + */ + public static double assignmentDistance(DistanceMetric metric, float[] a, float[] b) { + if (metric == DistanceMetric.COSINE) { + return VectorDistanceUtil.scalarCosineDistance(a, b, a.length); + } + return VectorDistanceUtil.scalarL2DistanceSquared(a, b, a.length); + } + + /** + * Returns a uniform sample of at most {@code sampleSize} vectors, by reservoir sampling. The + * sample skips null entries. {@code sampleSize} must be positive. + */ + public static List<float[]> reservoirSample(Iterator<float[]> iterator, int sampleSize, + Random random) { + if (sampleSize <= 0) { + throw new IllegalArgumentException("sampleSize must be > 0: " + sampleSize); + } + List<float[]> reservoir = new ArrayList<>(sampleSize); + long count = 0; + while (iterator.hasNext()) { + float[] v = iterator.next(); + if (v == null) { + continue; + } + if (count < sampleSize) { + reservoir.add(v); + } else { + long j = (long) (random.nextDouble() * (count + 1)); + if (j < sampleSize) { + reservoir.set((int) j, v); + } + } + count++; + } + return reservoir; + } + + /** + * Returns a copy of the vector scaled to unit length. If the norm is zero or NaN, it returns an + * unchanged copy. + */ + public static float[] l2Normalize(float[] v) { + double sumSq = 0.0; + for (float val : v) { + sumSq += (double) val * val; + } + double norm = Math.sqrt(sumSq); + if (norm == 0.0 || Double.isNaN(norm)) { + return v.clone(); + } + float[] out = new float[v.length]; + for (int i = 0; i < v.length; i++) { + out[i] = (float) (v[i] / norm); + } + return out; + } + + /** + * Selects {@code k} initial centroids by k-means++ seeding with the assignment distance. Under + * COSINE, each seed has unit length. + */ + static List<float[]> initializeKMeansPlusPlus(List<float[]> vectors, int k, DistanceMetric metric, + Random random) { + int n = vectors.size(); + List<float[]> centroids = new ArrayList<>(k); + centroids.add(seed(vectors.get(random.nextInt(n)), metric)); + double[] weight = new double[n]; + Arrays.fill(weight, Double.MAX_VALUE); + for (int c = 1; c < k; c++) { + float[] last = centroids.get(c - 1); + double total = 0.0; + for (int i = 0; i < n; i++) { + double d = assignmentDistance(metric, vectors.get(i), last); + // The L2 distance is already squared. Square the cosine distance, because k-means++ + // selects a seed with probability in proportion to the squared distance. + double w = metric == DistanceMetric.COSINE ? d * d : d; + if (w < weight[i]) { + weight[i] = w; + } + total += weight[i]; + } + int selected = n - 1; + if (total > 1e-12) { + double r = random.nextDouble() * total; + double accum = 0.0; + for (int i = 0; i < n; i++) { + accum += weight[i]; + if (accum >= r) { + selected = i; + break; + } + } + } else { + selected = random.nextInt(n); + } + centroids.add(seed(vectors.get(selected), metric)); + } + return centroids; + } + + /** + * Trains {@code k} centroids on the vector sample, then splits overloaded clusters if the + * configuration enables the split. All vectors must be non-null and have the same dimension. + * {@code k} must be positive and not more than the sample size. + */ + public static KMeansResult train(List<float[]> vectors, int k, KMeansConfig config) { + if (vectors == null || vectors.isEmpty()) { + throw new IllegalArgumentException("vectors list must not be null or empty"); + } + if (k <= 0) { + throw new IllegalArgumentException("k must be > 0: " + k); + } + if (k > vectors.size()) { + throw new IllegalArgumentException( + "k (" + k + ") cannot exceed number of vectors (" + vectors.size() + ")"); + } + int n = vectors.size(); + int dim = vectors.get(0).length; + for (int i = 0; i < n; i++) { + if (vectors.get(i) == null || vectors.get(i).length != dim) { + throw new IllegalArgumentException( + "Vector at index " + i + " is null or does not have" + " the expected dimension " + dim); + } + } + DistanceMetric metric = config.getDistanceMetric(); + Random random = + config.getRandomSeed() != null ? new Random(config.getRandomSeed()) : new Random(); + + List<float[]> centroids = initializeKMeansPlusPlus(vectors, k, metric, random); + Lloyd lloyd = new Lloyd(vectors, metric, random); + List<Double> distortionHistory = new ArrayList<>(); + double prevDistortion = Double.MAX_VALUE; + boolean converged = false; + int iteration; + for (iteration = 1; iteration <= config.getMaxIterations(); iteration++) { + double distortion = lloyd.iterate(centroids); + distortionHistory.add(distortion); + LOGGER.debug("K-means iteration {}: total distortion = {}", iteration, distortion); + if ( + distortion == 0.0 || (iteration > 1 && prevDistortion > 0.0 + && Math.abs(prevDistortion - distortion) / prevDistortion + < config.getConvergenceThreshold()) + ) { + converged = true; + break; + } + prevDistortion = distortion; + } + int iterations = Math.min(iteration, config.getMaxIterations()); + LOGGER.info("K-means finished after {} iterations, converged={}, distortion={}", iterations, + converged, lloyd.distortion); + + ClusterSkewMetrics preSplitSkew = ClusterSkewMetrics.compute(lloyd.clusterSizes); + ClusterSkewMetrics postSplitSkew = null; + if (config.isEnableSplitHeuristic() && n >= 2 * k) { + List<float[]> split = splitOverloaded(vectors, centroids, lloyd, + config.getSplitThresholdMultiplier() * n / k, metric, random); + if (split.size() > centroids.size()) { + centroids = split; + for (int pass = 0; pass < Math.max(1, config.getSplitRebalanceIterations()); pass++) { + lloyd.iterate(centroids); + } + postSplitSkew = ClusterSkewMetrics.compute(lloyd.clusterSizes); + LOGGER.info("Split heuristic raised centroid count from {} to {}; CV {} -> {}", k, + centroids.size(), preSplitSkew.getCoefficientOfVariation(), + postSplitSkew.getCoefficientOfVariation()); + } + } + ClusterSkewMetrics finalSkew = postSplitSkew != null ? postSplitSkew : preSplitSkew; + if (finalSkew.isSevereSkew()) { + LOGGER.warn("Severe cluster skew after training: CV = {}", + finalSkew.getCoefficientOfVariation()); + } + return new KMeansResult(centroids, k, iterations, converged, lloyd.distortion, + lloyd.clusterSizes, lloyd.assignments, finalSkew, preSplitSkew, postSplitSkew, dim, metric, + distortionHistory); + } + + public static KMeansResult train(List<float[]> vectors, int k) { + return train(vectors, k, KMeansConfig.defaultConfig()); + } + + /** + * Replaces each cluster with more members than {@code threshold} by two sub-centroids, which + * k-means++ seeds from its members. Returns {@code centroids} if no cluster is over the + * threshold. + */ + private static List<float[]> splitOverloaded(List<float[]> vectors, List<float[]> centroids, + Lloyd lloyd, double threshold, DistanceMetric metric, Random random) { + List<float[]> result = new ArrayList<>(centroids.size()); + boolean split = false; + for (int c = 0; c < centroids.size(); c++) { + if (lloyd.clusterSizes[c] > threshold) { + List<float[]> members = new ArrayList<>(lloyd.clusterSizes[c]); + for (int i = 0; i < vectors.size(); i++) { + if (lloyd.assignments[i] == c) { + members.add(vectors.get(i)); + } + } + result.addAll(initializeKMeansPlusPlus(members, 2, metric, random)); + split = true; + } else { + result.add(centroids.get(c)); + } + } + return split ? result : centroids; + } + + private static float[] seed(float[] v, DistanceMetric metric) { + return metric == DistanceMetric.COSINE ? l2Normalize(v) : v.clone(); + } + + /** + * State of the Lloyd iterations: the cluster of each vector, its distance and the cluster sizes. + * Each pass assigns vectors, reseeds empty clusters and updates the centroids. + */ + private static final class Lloyd { + private final List<float[]> vectors; + private final DistanceMetric metric; + private final Random random; + private final int[] assignments; + private final double[] distances; + private int[] clusterSizes; + private double distortion; + + Lloyd(List<float[]> vectors, DistanceMetric metric, Random random) { + this.vectors = vectors; + this.metric = metric; + this.random = random; + this.assignments = new int[vectors.size()]; + this.distances = new double[vectors.size()]; + } + + /** + * Runs one pass. It assigns each vector to its nearest centroid, reseeds empty clusters, and + * moves each centroid to the mean of its members. Returns the total distortion of the + * assignment. + */ + double iterate(List<float[]> centroids) { + int k = centroids.size(); + int n = vectors.size(); + int dim = vectors.get(0).length; + clusterSizes = new int[k]; + distortion = 0.0; + for (int i = 0; i < n; i++) { + float[] v = vectors.get(i); + int best = 0; + double bestDist = Double.MAX_VALUE; + for (int c = 0; c < k; c++) { + double d = assignmentDistance(metric, v, centroids.get(c)); + if (d < bestDist) { + bestDist = d; + best = c; + } + } + assignments[i] = best; + distances[i] = bestDist; + clusterSizes[best]++; + distortion += bestDist; + } + for (int c = 0; c < k; c++) { + if (clusterSizes[c] == 0) { + reseed(c); + } + } + double[][] sums = new double[k][dim]; + for (int i = 0; i < n; i++) { + float[] v = vectors.get(i); + double[] s = sums[assignments[i]]; + for (int d = 0; d < dim; d++) { + s[d] += v[d]; + } + } + for (int c = 0; c < k; c++) { + if (clusterSizes[c] == 0) { + continue; + } + float[] updated = new float[dim]; + for (int d = 0; d < dim; d++) { + updated[d] = (float) (sums[c][d] / clusterSizes[c]); + } + centroids.set(c, metric == DistanceMetric.COSINE ? l2Normalize(updated) : updated); + } + return distortion; + } + + /** + * Fills an empty cluster with the farthest vector of a cluster that has more than one member. + * If no such vector exists, it uses a random vector. The update step then puts the centroid on + * that vector. + */ + private void reseed(int c) { + int worst = -1; + double maxDist = -1.0; + for (int i = 0; i < assignments.length; i++) { + if (clusterSizes[assignments[i]] > 1 && distances[i] > maxDist) { + maxDist = distances[i]; + worst = i; + } + } + int idx = worst != -1 ? worst : random.nextInt(assignments.length); + LOGGER.debug("Cluster {} is empty; re-seeding with vector {}", c, idx); + clusterSizes[assignments[idx]]--; + assignments[idx] = c; + distances[idx] = 0.0; + clusterSizes[c]++; + } + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/VectorIndexTrainer.java b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/VectorIndexTrainer.java new file mode 100644 index 0000000000..408533fc0d --- /dev/null +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/index/vector/VectorIndexTrainer.java @@ -0,0 +1,213 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.index.vector; + +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.Random; +import org.apache.phoenix.cache.VectorCentroidCache; +import org.apache.phoenix.cache.VectorCentroidCache.CachedCentroids; +import org.apache.phoenix.jdbc.PhoenixConnection; +import org.apache.phoenix.optimize.DistanceMetric; +import org.apache.phoenix.schema.PColumn; +import org.apache.phoenix.schema.PTable; +import org.apache.phoenix.util.SchemaUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Trains the IVF centroids of an index, stores them, and records the generation in + * {@code SYSTEM.CATALOG}, for index creation and rebuild. + * <p> + * The training sample comes from a Bernoulli filter that the server evaluates on each data row. + * Reservoir sampling on the client then limits the sample size. Training runs k-means in memory. + */ +public final class VectorIndexTrainer { + + private static final Logger LOGGER = LoggerFactory.getLogger(VectorIndexTrainer.class); + /** + * Margin of the Bernoulli rate, so that the filter usually returns enough rows for the sample. + */ + private static final double OVERSAMPLE = 1.2; + + private VectorIndexTrainer() { + } + + /** + * Returns the vector column of the index. + * @throws IllegalStateException if the index has no vector column + */ + public static PColumn getIndexedVectorColumn(PTable index) { + for (PColumn column : index.getColumns()) { + if ( + column.getFamilyName() != null && column.getDataType() != null + && column.getDataType().isVectorType() + ) { + return column; + } + } + throw new IllegalStateException("Vector index " + index.getName() + " has no vector column"); + } + + /** + * Samples the data table and trains the number of centroids that the index definition requests. + * Returns null if the data table has too few vectors. + */ + public static KMeansResult train(PhoenixConnection conn, PTable dataTable, PTable index) + throws SQLException { + int lists = index.getVectorIvfLists(); + int sampleSize = index.getVectorIvfSampleSize(); + DistanceMetric metric = DistanceMetric.fromString(index.getVectorDistanceMetric()); + PColumn vectorColumn = getIndexedVectorColumn(index); + String expr = vectorColumn.getExpressionStr(); + String from = SchemaUtil.getEscapedFullTableName(dataTable.getName().getString()); + String notNull = " WHERE " + expr + " IS NOT NULL"; + + long count; + try (Statement stmt = conn.createStatement(); + ResultSet rs = stmt.executeQuery("SELECT COUNT(*) FROM " + from + notNull)) { + count = rs.next() ? rs.getLong(1) : 0; + } + if (count < lists) { + LOGGER.info("Deferring centroid training for vector index {}: {} non-null vectors, fewer than" + + " lists={}", index.getName(), count, lists); + return null; + } + double p = Math.min(1.0, OVERSAMPLE * sampleSize / count); + String sql = "SELECT " + expr + " FROM " + from + notNull; + List<float[]> sample = sample(conn, p < 1.0 ? sql + " AND RAND() < " + p : sql, sampleSize); + if (sample.size() < lists) { + // The Bernoulli filter can return fewer rows than lists. Then sample from a full scan. + sample = sample(conn, sql, sampleSize); + if (sample.size() < lists) { + LOGGER.info("Deferring centroid training for vector index {}: {} trainable vectors, fewer" + + " than lists={}", index.getName(), sample.size(), lists); + return null; + } + } + KMeansConfig config = KMeansConfig.newBuilder().distanceMetric(metric).build(); + LOGGER.info("Training vector index {} with {} sampled vectors, lists={}, metric={}", + index.getName(), sample.size(), lists, metric); + return KMeansTrainer.train(sample, lists, config); + } + + /** + * Trains a new centroid generation, stores its centroids and summary, and records it as the + * active generation in {@code SYSTEM.CATALOG}. It also puts the model in the centroid cache of + * this process. Returns the new generation ID, or null if the data table has too few vectors. + */ + public static Long trainAndRecord(PhoenixConnection conn, PTable dataTable, PTable index) + throws SQLException { + KMeansResult result = train(conn, dataTable, index); + if (result == null) { + return null; + } + String indexName = index.getName().getString(); + long generation = CentroidManager.nextGeneration(index.getVectorCentroidGeneration()); + CentroidManager.persistCentroids(conn, indexName, generation, result.getCentroids()); + CentroidManager.setGenerationAndLists(conn, index, generation, result.getEffectiveK()); + VectorCentroidCache.getInstance(conn.getQueryServices().getConfiguration()).put(indexName, + generation, new CachedCentroids(result.getCentroids(), result.getDistanceMetric())); + return generation; + } + + private static List<float[]> sample(PhoenixConnection conn, String sql, int sampleSize) + throws SQLException { + try (Statement stmt = conn.createStatement(); ResultSet rs = stmt.executeQuery(sql)) { + return KMeansTrainer.reservoirSample(new VectorIterator(rs), sampleSize, new Random()); + } catch (VectorIterator.ReadException e) { + throw e.getCause(); + } + } + + /** + * Iterates the vectors of a result set as float arrays. It returns null for a null or non-finite + * vector, and it wraps an SQLException in a ReadException. + */ + private static final class VectorIterator implements Iterator<float[]> { + private final ResultSet rs; + private Boolean hasNext; + + VectorIterator(ResultSet rs) { + this.rs = rs; + } + + @Override + public boolean hasNext() { + if (hasNext == null) { + try { + hasNext = rs.next(); + } catch (SQLException e) { + throw new ReadException(e); + } + } + return hasNext; + } + + @Override + public float[] next() { + if (!hasNext()) { + throw new NoSuchElementException(); + } + hasNext = null; + try { + float[] vector = toFloats(rs.getObject(1)); + for (int i = 0; vector != null && i < vector.length; i++) { + if (!Float.isFinite(vector[i])) { + // K-means cannot place a non-finite vector, for example a DOUBLE element beyond + // float range that narrows to infinity. The sample skips it, as it skips a null one. + return null; + } + } + return vector; + } catch (SQLException e) { + throw new ReadException(e); + } + } + + private static float[] toFloats(Object value) { + if (value instanceof float[]) { + return (float[]) value; + } + if (value instanceof double[]) { + double[] d = (double[]) value; + float[] f = new float[d.length]; + for (int i = 0; i < d.length; i++) { + f[i] = (float) d[i]; + } + return f; + } + return null; + } + + private static final class ReadException extends RuntimeException { + ReadException(SQLException cause) { + super(cause); + } + + @Override + public synchronized SQLException getCause() { + return (SQLException) super.getCause(); + } + } + } +} diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/QueryOptimizer.java b/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/QueryOptimizer.java index 0300227a7d..a2b96e2ece 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/QueryOptimizer.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/optimize/QueryOptimizer.java @@ -293,7 +293,12 @@ public class QueryOptimizer { } } - List<PTable> indexes = Lists.newArrayList(dataPlan.getTableRef().getTable().getIndexes()); + List<PTable> indexes = Lists.newArrayList(); + for (PTable idx : dataPlan.getTableRef().getTable().getIndexes()) { + if (!idx.isVectorIndex()) { + indexes.add(idx); + } + } if ( dataPlan.isApplicable() && (indexes.isEmpty() || dataPlan.isDegenerate() || dataPlan.getTableRef().hasDynamicCols() || select.getHint().hasHint(Hint.NO_INDEX)) diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/query/QueryServices.java b/phoenix-core-client/src/main/java/org/apache/phoenix/query/QueryServices.java index 085ac34a64..0f194aa89a 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/query/QueryServices.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/query/QueryServices.java @@ -150,6 +150,9 @@ public interface QueryServices extends SQLCloseable { "phoenix.coprocessor.maxMetaDataCacheSize"; public static final String MAX_CLIENT_METADATA_CACHE_SIZE_ATTRIB = "phoenix.client.maxMetaDataCacheSize"; + // Maximum bytes of centroid vectors that each process caches, for all vector index generations + public static final String VECTOR_CENTROID_CACHE_MAX_BYTES_ATTRIB = + "phoenix.vector.centroid.cache.maxBytes"; public static final String HA_GROUP_NAME_ATTRIB = "phoenix.ha.group"; public static final String AUTO_UPGRADE_WHITELIST_ATTRIB = "phoenix.client.autoUpgradeWhiteList"; // Mainly for testing to force spilling diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/query/QueryServicesOptions.java b/phoenix-core-client/src/main/java/org/apache/phoenix/query/QueryServicesOptions.java index 4e3c29b6c3..dd801cb6c4 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/query/QueryServicesOptions.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/query/QueryServicesOptions.java @@ -232,6 +232,7 @@ public class QueryServicesOptions { // mins public static final long DEFAULT_MAX_SERVER_METADATA_CACHE_SIZE = 1024L * 1024L * 20L; // 20 Mb public static final long DEFAULT_MAX_CLIENT_METADATA_CACHE_SIZE = 1024L * 1024L * 10L; // 10 Mb + public static final long DEFAULT_VECTOR_CENTROID_CACHE_MAX_BYTES = 1024L * 1024L * 256L; // 256 Mb public static final int DEFAULT_GROUPBY_ESTIMATED_DISTINCT_VALUES = 1000; public static final int DEFAULT_CLOCK_SKEW_INTERVAL = 2000; public static final boolean DEFAULT_INDEX_FAILURE_HANDLING_REBUILD = true; // auto rebuild on diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/schema/MetaDataClient.java b/phoenix-core-client/src/main/java/org/apache/phoenix/schema/MetaDataClient.java index 0f3368a234..26ce747202 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/schema/MetaDataClient.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/schema/MetaDataClient.java @@ -201,6 +201,7 @@ import org.apache.hadoop.hbase.security.access.Permission; import org.apache.hadoop.hbase.util.Bytes; import org.apache.hadoop.hbase.util.Pair; import org.apache.hadoop.util.StringUtils; +import org.apache.phoenix.cache.VectorCentroidCache; import org.apache.phoenix.compile.ColumnResolver; import org.apache.phoenix.compile.FromCompiler; import org.apache.phoenix.compile.IndexExpressionCompiler; @@ -227,6 +228,8 @@ import org.apache.phoenix.expression.function.PartitionIdFunction; import org.apache.phoenix.expression.function.PhoenixRowTimestampFunction; import org.apache.phoenix.hbase.index.covered.update.ColumnReference; import org.apache.phoenix.index.IndexMaintainer; +import org.apache.phoenix.index.vector.CentroidManager; +import org.apache.phoenix.index.vector.VectorIndexTrainer; import org.apache.phoenix.jdbc.PhoenixConnection; import org.apache.phoenix.jdbc.PhoenixDatabaseMetaData; import org.apache.phoenix.jdbc.PhoenixStatement; @@ -2006,11 +2009,12 @@ public class MetaDataClient { // Vector indexes remain in BUILDING state until centroids are initialized and assigned. if (table.isVectorIndex()) { + MutationState state = trainVectorIndex(table, dataTable); if (ValidateLastDDLTimestampUtil.getValidateLastDdlTimestampEnabled(connection)) { connection.removeTable(connection.getTenantId(), dataTable.getName().getString(), null, dataTable.getTimeStamp()); } - return new MutationState(0, 0, connection); + return state; } // If our connection is at a fixed point-in-time, we need to open a new @@ -2029,6 +2033,20 @@ public class MetaDataClient { return state; } + /** + * Trains and records the initial centroid generation for a vector index, retaining BUILDING + * state. Training is deferred if the data table contains fewer non-null vectors than requested + * lists. + */ + private MutationState trainVectorIndex(PTable index, PTable dataTable) throws SQLException { + try (PhoenixConnection internal = CentroidManager.newInternalConnection(connection)) { + VectorIndexTrainer.trainAndRecord(internal, dataTable, index); + } + connection.removeTable(connection.getTenantId(), index.getName().getString(), + dataTable.getName().getString(), HConstants.LATEST_TIMESTAMP); + return new MutationState(0, 0, connection); + } + public MutationState createCDC(CreateCDCStatement statement) throws SQLException { ColumnResolver resolver = FromCompiler.getResolver(NamedTableNode.create(statement.getDataTable()), connection); @@ -4634,27 +4652,10 @@ public class MetaDataClient { return; } String indexName = index.getName().getString(); - try { - Properties props = new Properties(); - Properties clientInfo = connection.getClientInfo(); - for (String key : clientInfo.stringPropertyNames()) { - if ( - !PhoenixRuntime.TENANT_ID_ATTRIB.equals(key) - && !PhoenixRuntime.CURRENT_SCN_ATTRIB.equals(key) - ) { - props.setProperty(key, clientInfo.getProperty(key)); - } - } - try ( - PhoenixConnection conn = - new PhoenixConnection(connection, connection.getQueryServices(), props); - PreparedStatement ps = - conn.prepareStatement("DELETE FROM " + PhoenixDatabaseMetaData.SYSTEM_VECTOR_CENTROID_NAME - + " WHERE " + PhoenixDatabaseMetaData.INDEX_NAME + " = ?")) { - conn.setAutoCommit(true); - ps.setString(1, indexName); - ps.executeUpdate(); - } + VectorCentroidCache.getInstance(connection.getQueryServices().getConfiguration()) + .invalidate(indexName); + try (PhoenixConnection internal = CentroidManager.newInternalConnection(connection)) { + CentroidManager.deleteAllCentroids(internal, indexName); } catch (SQLException e) { LOGGER.warn("Could not delete the centroids of dropped vector index {}; its " + PhoenixDatabaseMetaData.SYSTEM_VECTOR_CENTROID_NAME + " rows remain.", indexName, e); diff --git a/phoenix-core-client/src/main/java/org/apache/phoenix/util/IndexUtil.java b/phoenix-core-client/src/main/java/org/apache/phoenix/util/IndexUtil.java index 830c2d49c6..7330e537bf 100644 --- a/phoenix-core-client/src/main/java/org/apache/phoenix/util/IndexUtil.java +++ b/phoenix-core-client/src/main/java/org/apache/phoenix/util/IndexUtil.java @@ -174,7 +174,7 @@ public class IndexUtil { // row key was already done, so here we just need to convert from one built-in type to // another. public static PDataType getIndexColumnDataType(boolean isNullable, PDataType dataType) { - if (dataType == null || !isNullable || !dataType.isFixedWidth()) { + if (dataType == null || !isNullable || !dataType.isFixedWidth() || dataType.isVectorType()) { return dataType; } // for fixed length numeric types and boolean diff --git a/phoenix-core-server/src/main/java/org/apache/phoenix/coprocessor/MetaDataEndpointImpl.java b/phoenix-core-server/src/main/java/org/apache/phoenix/coprocessor/MetaDataEndpointImpl.java index 5b42865d20..96d27b0706 100644 --- a/phoenix-core-server/src/main/java/org/apache/phoenix/coprocessor/MetaDataEndpointImpl.java +++ b/phoenix-core-server/src/main/java/org/apache/phoenix/coprocessor/MetaDataEndpointImpl.java @@ -4608,7 +4608,16 @@ public class MetaDataEndpointImpl extends MetaDataProtocol implements RegionCopr int disableTimeStampKVIndex = -1; int indexStateKVIndex = 0; int index = 0; + // Find if this mutation records an active or building vector centroid generation + boolean setsVectorGeneration = false; for (Cell cell : newKVs) { + if ( + Bytes.compareTo(cell.getQualifierArray(), cell.getQualifierOffset(), + cell.getQualifierLength(), VECTOR_CENTROID_GENERATION_BYTES, 0, + VECTOR_CENTROID_GENERATION_BYTES.length) == 0 + ) { + setsVectorGeneration = true; + } if ( Bytes.compareTo(cell.getQualifierArray(), cell.getQualifierOffset(), cell.getQualifierLength(), INDEX_STATE_BYTES, 0, INDEX_STATE_BYTES.length) == 0 @@ -4677,6 +4686,12 @@ public class MetaDataEndpointImpl extends MetaDataProtocol implements RegionCopr PIndexState currentState = PIndexState .fromSerializedValue(currentStateKV.getValueArray()[currentStateKV.getValueOffset()]); + if (setsVectorGeneration && newState != currentState) { + // A generation record must not change the index state, so keep the current state + newState = currentState; + newKVs.set(indexStateKVIndex, PhoenixKeyValueUtil.newKeyValue(key, TABLE_FAMILY_BYTES, + INDEX_STATE_BYTES, timeStamp, Bytes.toBytes(newState.getSerializedValue()))); + } // Timestamp of INDEX_STATE gets updated with each call long actualTimestamp = currentStateKV.getTimestamp(); long curTimeStampVal = 0; @@ -4828,7 +4843,7 @@ public class MetaDataEndpointImpl extends MetaDataProtocol implements RegionCopr } PTable returnTable = null; - if (currentState != newState || disableTimeStampKVIndex != -1) { + if (currentState != newState || disableTimeStampKVIndex != -1 || setsVectorGeneration) { // make a copy of tableMetadata so we can add to it tableMetadata = new ArrayList<Mutation>(tableMetadata); // Always include the empty column value at latest timestamp so @@ -4874,6 +4889,7 @@ public class MetaDataEndpointImpl extends MetaDataProtocol implements RegionCopr if ( setRowKeyOrderOptimizableCell || disableTimeStampKVIndex != -1 || currentState.isDisabled() || newState == PIndexState.BUILDING + || setsVectorGeneration ) { returnTable = doGetTable(tenantId, schemaName, tableName, HConstants.LATEST_TIMESTAMP, rowLock, request.getClientVersion()); diff --git a/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorCentroidTableIT.java b/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorCentroidTableIT.java index 1192fe3fca..91b389059a 100644 --- a/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorCentroidTableIT.java +++ b/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorCentroidTableIT.java @@ -40,7 +40,9 @@ import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.Statement; import java.sql.Types; +import java.util.Arrays; import java.util.List; +import org.apache.phoenix.index.vector.CentroidManager; import org.apache.phoenix.jdbc.PhoenixConnection; import org.apache.phoenix.schema.PColumn; import org.apache.phoenix.schema.PTable; @@ -218,4 +220,102 @@ public class VectorCentroidTableIT extends ParallelStatsDisabledIT { } } } + + private static void assertCentroidsEqual(List<float[]> expected, List<float[]> actual) { + assertEquals(expected.size(), actual.size()); + for (int i = 0; i < expected.size(); i++) { + assertArrayEquals(expected.get(i), actual.get(i), 0f); + } + } + + @Test + public void testPersistAndLoadRoundTrip() throws Exception { + String indexName = "TEST_VECTOR_IDX_" + generateUniqueName(); + List<float[]> centroids = + Arrays.asList(new float[] { 1, 2 }, new float[] { 3, 4 }, new float[] { 5, 6 }); + try (Connection conn = DriverManager.getConnection(getUrl())) { + CentroidManager.persistCentroids(conn, indexName, 1L, centroids); + assertCentroidsEqual(centroids, CentroidManager.loadCentroids(conn, indexName, 1L)); + } + } + + @Test + public void testGenerationIsolation() throws Exception { + String indexName = "TEST_VECTOR_IDX_" + generateUniqueName(); + List<float[]> gen1 = Arrays.asList(new float[] { 1, 1 }, new float[] { 2, 2 }); + List<float[]> gen2 = Arrays.asList(new float[] { 9, 9 }); + try (Connection conn = DriverManager.getConnection(getUrl())) { + CentroidManager.persistCentroids(conn, indexName, 1L, gen1); + CentroidManager.persistCentroids(conn, indexName, 2L, gen2); + assertCentroidsEqual(gen1, CentroidManager.loadCentroids(conn, indexName, 1L)); + assertCentroidsEqual(gen2, CentroidManager.loadCentroids(conn, indexName, 2L)); + assertTrue(CentroidManager.loadCentroids(conn, indexName, 3L).isEmpty()); + } + } + + /** Makes sure that index names that differ only in case keep separate centroid rows. */ + @Test + public void testIndexNamesAreCaseSensitive() throws Exception { + String upper = "S.IDX_" + generateUniqueName(); + String mixed = upper.replace("IDX_", "idx_"); + try (Connection conn = DriverManager.getConnection(getUrl())) { + CentroidManager.persistCentroids(conn, upper, 1L, Arrays.asList(new float[] { 1 })); + CentroidManager.persistCentroids(conn, mixed, 1L, Arrays.asList(new float[] { 2 })); + assertArrayEquals(new float[] { 1 }, CentroidManager.loadCentroids(conn, upper, 1L).get(0), + 0f); + assertArrayEquals(new float[] { 2 }, CentroidManager.loadCentroids(conn, mixed, 1L).get(0), + 0f); + } + } + + @Test + public void testDeleteGenerationAndAll() throws Exception { + String indexName = "TEST_VECTOR_IDX_" + generateUniqueName(); + String other = "TEST_VECTOR_IDX_" + generateUniqueName(); + try (Connection conn = DriverManager.getConnection(getUrl())) { + CentroidManager.persistCentroids(conn, indexName, 1L, Arrays.asList(new float[] { 1 })); + CentroidManager.persistCentroids(conn, indexName, 2L, Arrays.asList(new float[] { 2 })); + CentroidManager.persistCentroids(conn, other, 1L, Arrays.asList(new float[] { 3 })); + + CentroidManager.deleteGeneration(conn, indexName, 1L); + assertTrue(CentroidManager.loadCentroids(conn, indexName, 1L).isEmpty()); + assertEquals(1, CentroidManager.loadCentroids(conn, indexName, 2L).size()); + + CentroidManager.deleteAllCentroids(conn, indexName); + assertTrue(CentroidManager.loadCentroids(conn, indexName, 2L).isEmpty()); + assertEquals(1, CentroidManager.loadCentroids(conn, other, 1L).size()); + assertFalse("Deletes restore the caller's autocommit setting", conn.getAutoCommit()); + } + } + + /** + * Makes sure that a generation record through the metadata endpoint keeps the index state, and + * advances the DDL timestamp so that clients refresh cached metadata. + */ + @Test + public void testSetGenerationAndListsThroughEndpoint() throws Exception { + String dataTable = "DATA_" + generateUniqueName(); + String vectorIdx = "IDX_" + generateUniqueName(); + try (Connection conn = DriverManager.getConnection(getUrl())) { + try (Statement stmt = conn.createStatement()) { + stmt.execute("CREATE TABLE " + dataTable + " (ID VARCHAR PRIMARY KEY, V VECTOR(FLOAT, 4))"); + stmt.execute("CREATE VECTOR INDEX " + vectorIdx + " ON " + dataTable + + " (V) WITH (algorithm = 'IVF', metric = 'L2', lists = 2, sample_size = 10)"); + } + PhoenixConnection pconn = conn.unwrap(PhoenixConnection.class); + PTable before = pconn.getTableNoCache(vectorIdx); + assertEquals("An empty table defers training", null, before.getVectorCentroidGeneration()); + + long generation = CentroidManager.nextGeneration(null); + try (PhoenixConnection internal = CentroidManager.newInternalConnection(pconn)) { + CentroidManager.setGenerationAndLists(internal, before, generation, 3); + } + PTable after = pconn.getTableNoCache(vectorIdx); + assertEquals(Long.valueOf(generation), after.getVectorCentroidGeneration()); + assertEquals(Integer.valueOf(3), after.getVectorIvfLists()); + assertEquals(before.getIndexState(), after.getIndexState()); + assertTrue("The index DDL timestamp advances", + after.getLastDDLTimestamp() > before.getLastDDLTimestamp()); + } + } } diff --git a/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorIndexIT.java b/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorIndexIT.java index b8d5e97f47..05a11f4ac2 100644 --- a/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorIndexIT.java +++ b/phoenix-core/src/it/java/org/apache/phoenix/end2end/VectorIndexIT.java @@ -45,6 +45,7 @@ import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import java.io.IOException; +import java.sql.Array; import java.sql.Connection; import java.sql.DriverManager; import java.sql.PreparedStatement; @@ -925,4 +926,160 @@ public class VectorIndexIT extends ParallelStatsDisabledIT { } } } + + private static void loadClusteredVectors(Connection conn, String tableName, int rows) + throws SQLException { + try (PreparedStatement ps = + conn.prepareStatement("UPSERT INTO " + tableName + " (ID, V, LABEL) VALUES (?, ?, ?)")) { + for (int i = 0; i < rows; i++) { + ps.setString(1, String.format("id_%03d", i)); + Float[] vec = new Float[] { (float) (i % 4), (float) ((i + 1) % 4), (float) ((i + 2) % 4), + (float) ((i + 3) % 4) }; + Array array = conn.createArrayOf("FLOAT", vec); + ps.setArray(2, array); + ps.setString(3, "label_" + i); + ps.executeUpdate(); + } + } + conn.commit(); + } + + private static int countCentroids(Connection conn, String indexName, Long generation) + throws SQLException { + String sql = "SELECT COUNT(*) FROM " + SYSTEM_VECTOR_CENTROID_NAME + " WHERE " + INDEX_NAME + + " = ?" + (generation == null ? "" : " AND " + GENERATION_ID + " = ?"); + try (PreparedStatement ps = conn.prepareStatement(sql)) { + ps.setString(1, indexName); + if (generation != null) { + ps.setLong(2, generation); + } + try (ResultSet rs = ps.executeQuery()) { + assertTrue(rs.next()); + return rs.getInt(1); + } + } + } + + /** + * Verifies that synchronous index creation trains initial centroids, records active generation + * metadata, and retains BUILDING state pending population. + */ + @Test + public void testSynchronousVectorIndexTraining() throws Exception { + String tableName = "T_VEC_POP_" + generateUniqueName(); + String indexName = "IDX_VEC_POP_" + generateUniqueName(); + try (Connection conn = DriverManager.getConnection(getUrl())) { + try (Statement stmt = conn.createStatement()) { + stmt.execute("CREATE TABLE " + tableName + + " (ID VARCHAR NOT NULL PRIMARY KEY, V VECTOR(FLOAT, 4), LABEL VARCHAR)"); + } + loadClusteredVectors(conn, tableName, 100); + long before = System.currentTimeMillis(); + try (Statement stmt = conn.createStatement()) { + stmt.execute( + "CREATE VECTOR INDEX " + indexName + " ON " + tableName + " (V) " + "INCLUDE (LABEL) " + + "WITH (algorithm = 'IVF', metric = 'L2', lists = 4, sample_size = 100)"); + } + PTable index = conn.unwrap(PhoenixConnection.class).getTableNoCache(indexName); + assertEquals(PIndexState.BUILDING, index.getIndexState()); + Long generation = index.getVectorCentroidGeneration(); + assertNotNull("Training records a generation", generation); + assertTrue("The first generation is the training time", generation >= before); + assertTrue(index.getVectorIvfLists() >= 4); + assertEquals(index.getVectorIvfLists().intValue(), + countCentroids(conn, indexName, generation)); + } + } + + /** + * Verifies that centroid training is deferred when the table contains fewer vectors than + * requested lists. + */ + @Test + public void testTrainingDeferredOnEmptyTable() throws Exception { + String tableName = "T_VEC_" + generateUniqueName(); + String indexName = "IDX_VEC_" + generateUniqueName(); + try (Connection conn = DriverManager.getConnection(getUrl())) { + try (Statement stmt = conn.createStatement()) { + stmt.execute("CREATE TABLE " + tableName + + " (ID VARCHAR NOT NULL PRIMARY KEY, V VECTOR(FLOAT, 4), LABEL VARCHAR)"); + stmt.execute("CREATE VECTOR INDEX " + indexName + " ON " + tableName + + " (V) WITH (algorithm = 'IVF', metric = 'L2', lists = 4, sample_size = 100)"); + } + PTable index = conn.unwrap(PhoenixConnection.class).getTableNoCache(indexName); + assertEquals(PIndexState.BUILDING, index.getIndexState()); + assertNull(index.getVectorCentroidGeneration()); + assertEquals(0, countCentroids(conn, indexName, null)); + } + } + + /** + * Makes sure that the centroid rows of a quoted index name use its exact case. DROP INDEX must + * delete the rows, and a re-created index must get a higher generation ID. The higher ID prevents + * reuse of a stale cached model. + */ + @Test + public void testQuotedIndexNameTrainDropRecreate() throws Exception { + String schema = generateUniqueName(); + String tableName = schema + ".T_" + generateUniqueName(); + String indexName = "myIdx"; + String fullIndexName = schema + "." + indexName; + try (Connection conn = DriverManager.getConnection(getUrl())) { + try (Statement stmt = conn.createStatement()) { + stmt.execute("CREATE TABLE " + tableName + + " (ID VARCHAR NOT NULL PRIMARY KEY, V VECTOR(FLOAT, 4), LABEL VARCHAR)"); + } + loadClusteredVectors(conn, tableName, 40); + String ddl = "CREATE VECTOR INDEX \"" + indexName + "\" ON " + tableName + + " (V) WITH (algorithm = 'IVF', metric = 'L2', lists = 2, sample_size = 40)"; + try (Statement stmt = conn.createStatement()) { + stmt.execute(ddl); + } + PhoenixConnection pconn = conn.unwrap(PhoenixConnection.class); + PTable index = pconn.getTableNoCache(fullIndexName); + assertEquals(fullIndexName, index.getName().getString()); + Long firstGeneration = index.getVectorCentroidGeneration(); + assertNotNull(firstGeneration); + assertTrue(countCentroids(conn, fullIndexName, firstGeneration) >= 2); + assertEquals(0, countCentroids(conn, fullIndexName.toUpperCase(), null)); + + try (Statement stmt = conn.createStatement()) { + stmt.execute("DROP INDEX \"" + indexName + "\" ON " + tableName); + } + assertEquals(0, countCentroids(conn, fullIndexName, null)); + + try (Statement stmt = conn.createStatement()) { + stmt.execute(ddl); + } + Long secondGeneration = pconn.getTableNoCache(fullIndexName).getVectorCentroidGeneration(); + assertNotNull(secondGeneration); + assertTrue(secondGeneration > firstGeneration); + } + } + + /** + * Makes sure that the server evaluates the RAND() sample filter on each row, so that the sample + * keeps about the requested fraction of rows. + */ + @Test + public void testTrainingSampleFilterIsEvaluatedPerRowOnServer() throws Exception { + String tableName = "T_VEC_" + generateUniqueName(); + try (Connection conn = DriverManager.getConnection(getUrl())) { + try (Statement stmt = conn.createStatement()) { + stmt.execute("CREATE TABLE " + tableName + + " (ID VARCHAR NOT NULL PRIMARY KEY, V VECTOR(FLOAT, 4), LABEL VARCHAR)"); + } + loadClusteredVectors(conn, tableName, 1000); + String sql = "SELECT V FROM " + tableName + " WHERE V IS NOT NULL AND RAND() < 0.3"; + String plan = QueryUtil.getExplainPlan(conn.createStatement().executeQuery("EXPLAIN " + sql)); + assertTrue(plan, plan.contains("SERVER FILTER BY") && plan.contains("RAND()")); + int rows = 0; + try (ResultSet rs = conn.createStatement().executeQuery(sql)) { + while (rs.next()) { + rows++; + } + } + assertTrue("Sampled " + rows + " of 1000 at p=0.3", rows > 200 && rows < 400); + } + } } diff --git a/phoenix-core/src/test/java/org/apache/phoenix/cache/VectorCentroidCacheTest.java b/phoenix-core/src/test/java/org/apache/phoenix/cache/VectorCentroidCacheTest.java new file mode 100644 index 0000000000..eb234274e7 --- /dev/null +++ b/phoenix-core/src/test/java/org/apache/phoenix/cache/VectorCentroidCacheTest.java @@ -0,0 +1,315 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.cache; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.sql.Connection; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.phoenix.cache.VectorCentroidCache.CachedCentroids; +import org.apache.phoenix.optimize.DistanceMetric; +import org.apache.phoenix.schema.types.PVectorDouble; +import org.apache.phoenix.schema.types.PVectorFloat; +import org.junit.Test; +import org.mockito.stubbing.Answer; + +public class VectorCentroidCacheTest { + + private static final long MAX_BYTES = 1024L * 1024L; + + private static List<float[]> centroids(float[]... c) { + return Arrays.asList(c); + } + + /** Returns a mock connection that serves the given centroid rows and counts each load. */ + private static Connection connectionReturning(List<float[]> rows, AtomicInteger loads) + throws SQLException { + Connection conn = mock(Connection.class); + PreparedStatement ps = mock(PreparedStatement.class); + when(conn.prepareStatement(anyString())).thenReturn(ps); + when(ps.executeQuery()).thenAnswer((Answer<ResultSet>) invocation -> { + loads.incrementAndGet(); + ResultSet rs = mock(ResultSet.class); + int[] pos = { -1 }; + when(rs.next()).thenAnswer(i -> ++pos[0] < rows.size()); + when(rs.getBytes(1)).thenAnswer(i -> PVectorFloat.INSTANCE.toBytes(rows.get(pos[0]))); + return rs; + }); + return conn; + } + + @Test + public void testAssignNearestCentroid() { + CachedCentroids c = new CachedCentroids( + centroids(new float[] { 0, 0 }, new float[] { 10, 0 }, new float[] { 0, 10 }), + DistanceMetric.L2); + assertEquals(1, c.assign(new float[] { 9.5f, 0.5f })); + assertEquals(2, c.assign(new float[] { 0.5f, 9f })); + assertEquals(0, c.assign(new float[] { 1f, 1f })); + } + + @Test + public void testNearestNInDistanceOrder() { + CachedCentroids c = new CachedCentroids(centroids(new float[] { 0, 0 }, new float[] { 10, 0 }, + new float[] { 0, 10 }, new float[] { 20, 20 }), DistanceMetric.L2); + assertArrayEquals(new int[] { 1, 0 }, c.nearest(new float[] { 9.5f, 0.5f }, 2)); + assertArrayEquals(new int[] { 1, 0, 2, 3 }, c.nearest(new float[] { 9.5f, 0.5f }, 10)); + } + + /** Verifies that an INNER_PRODUCT index assigns a vector to a centroid by L2 distance. */ + @Test + public void testInnerProductAssignsByL2() { + CachedCentroids c = new CachedCentroids(centroids(new float[] { 9, 0 }, new float[] { 100, 0 }), + DistanceMetric.INNER_PRODUCT); + assertEquals(0, c.assign(new float[] { 10, 0 })); + } + + @Test + public void testCosineAssignsByAngle() { + CachedCentroids c = new CachedCentroids(centroids(new float[] { 1, 0 }, new float[] { 0, 1 }), + DistanceMetric.COSINE); + assertEquals(0, c.assign(new float[] { 100, 1 })); + assertEquals(1, c.assign(new float[] { 0.01f, 5 })); + } + + /** Verifies that packed FLOAT and DOUBLE forms of one vector go to the same centroid. */ + @Test + public void testAssignPackedFloatAndDoubleVectors() { + CachedCentroids c = new CachedCentroids( + centroids(new float[] { 0, 0, 0 }, new float[] { 5, 5, 5 }), DistanceMetric.L2); + byte[] f = PVectorFloat.INSTANCE.toBytes(new float[] { 4, 4, 4 }); + byte[] d = PVectorDouble.INSTANCE.toBytes(new double[] { 4, 4, 4 }); + assertEquals(1, c.assign(f, 0, f.length, PVectorFloat.INSTANCE)); + assertEquals(1, c.assign(d, 0, d.length, PVectorDouble.INSTANCE)); + } + + /** + * Verifies that a vector with a DOUBLE element beyond float range routes by its direction. The + * element narrows to infinity, and assignment uses the largest float of the same sign instead. + * Without saturation, the distances are infinite or NaN, and the tie decides the centroid. + */ + @Test + public void testOutOfFloatRangeElementRoutesByDirection() { + CachedCentroids cosine = new CachedCentroids( + centroids(new float[] { 1, 0 }, new float[] { 0, 1 }), DistanceMetric.COSINE); + byte[] d = PVectorDouble.INSTANCE.toBytes(new double[] { 1, 1e300 }); + assertEquals(1, cosine.assign(d, 0, d.length, PVectorDouble.INSTANCE)); + // nearest() also saturates the query. Without saturation, both distances are NaN and tie, + // and the tie puts ID 0 first. + assertArrayEquals(new int[] { 1, 0 }, + cosine.nearest(new float[] { 1, Float.POSITIVE_INFINITY }, 2)); + + CachedCentroids l2 = new CachedCentroids( + centroids(new float[] { 0, 0 }, new float[] { -1e37f, 0 }), DistanceMetric.L2); + assertEquals(1, l2.assign(new float[] { Float.NEGATIVE_INFINITY, 0 })); + } + + /** + * Verifies that nearest() breaks distance ties by ascending centroid ID, as assign() does. Under + * L2, a saturated component dominates the distance to every ordinary centroid, and the distances + * tie. A probe for such a vector must start at the list that holds its rows. + */ + @Test + public void testNearestBreaksTiesLikeAssign() { + CachedCentroids l2 = new CachedCentroids( + centroids(new float[] { 0, 0 }, new float[] { 1, 1 }, new float[] { 2, 2 }), + DistanceMetric.L2); + float[] v = { Float.POSITIVE_INFINITY, 0 }; + assertEquals(0, l2.assign(v)); + assertArrayEquals(new int[] { 0 }, l2.nearest(v, 1)); + assertArrayEquals(new int[] { 0, 1 }, l2.nearest(v, 2)); + assertArrayEquals(new int[] { 0, 1, 2 }, l2.nearest(v, 3)); + } + + @Test + public void testDimensionMismatchRejected() { + CachedCentroids c = new CachedCentroids(centroids(new float[] { 0, 0 }), DistanceMetric.L2); + try { + c.assign(new float[] { 1, 2, 3 }); + fail("Expected a dimension mismatch"); + } catch (IllegalArgumentException expected) { + } + } + + /** Verifies that the cache keeps more than one centroid generation of an index. */ + @Test + public void testGenerationsCoexist() { + VectorCentroidCache cache = new VectorCentroidCache(MAX_BYTES); + CachedCentroids g1 = new CachedCentroids(centroids(new float[] { 1 }), DistanceMetric.L2); + CachedCentroids g2 = new CachedCentroids(centroids(new float[] { 2 }), DistanceMetric.L2); + cache.put("S.IDX", 100L, g1); + cache.put("S.IDX", 200L, g2); + assertSame(g1, cache.getIfPresent("S.IDX", 100L)); + assertSame(g2, cache.getIfPresent("S.IDX", 200L)); + assertNull(cache.getIfPresent("S.IDX", 300L)); + } + + /** Verifies that the cache compares index names with case sensitivity. */ + @Test + public void testNamesAreCaseSensitive() { + VectorCentroidCache cache = new VectorCentroidCache(MAX_BYTES); + cache.put("S.myIdx", 1L, new CachedCentroids(centroids(new float[] { 1 }), DistanceMetric.L2)); + assertNull(cache.getIfPresent("S.MYIDX", 1L)); + } + + @Test + public void testLoadOnMissThenHit() throws Exception { + VectorCentroidCache cache = new VectorCentroidCache(MAX_BYTES); + AtomicInteger loads = new AtomicInteger(); + Connection conn = + connectionReturning(centroids(new float[] { 0, 0 }, new float[] { 10, 10 }), loads); + CachedCentroids first = cache.get(conn, "IDX", 7L, DistanceMetric.L2); + assertEquals(2, first.size()); + assertArrayEquals(new float[] { 10, 10 }, first.getCentroid(1), 0f); + assertSame(first, cache.get(conn, "IDX", 7L, DistanceMetric.L2)); + assertEquals(1, loads.get()); + } + + /** + * Verifies that a load that finds no centroids fails and puts no entry in the cache. A later load + * for the same key can then succeed. + */ + @Test + public void testEmptyLoadIsNotCached() throws Exception { + VectorCentroidCache cache = new VectorCentroidCache(MAX_BYTES); + AtomicInteger loads = new AtomicInteger(); + try { + cache.get(connectionReturning(new ArrayList<>(), loads), "IDX", 1L, DistanceMetric.L2); + fail("Expected an empty generation to fail"); + } catch (SQLException expected) { + } + assertNull(cache.getIfPresent("IDX", 1L)); + assertNotNull(cache.get(connectionReturning(centroids(new float[] { 1 }), loads), "IDX", 1L, + DistanceMetric.L2)); + } + + /** + * Verifies that concurrent misses for the same key share one load. The load blocks until every + * caller waits inside the cache, so all callers miss at the same time. + */ + @Test(timeout = 30000) + public void testConcurrentMissesLoadOnce() throws Exception { + VectorCentroidCache cache = new VectorCentroidCache(MAX_BYTES); + AtomicInteger loads = new AtomicInteger(); + CountDownLatch release = new CountDownLatch(1); + Connection conn = mock(Connection.class); + PreparedStatement ps = mock(PreparedStatement.class); + when(conn.prepareStatement(anyString())).thenReturn(ps); + when(ps.executeQuery()).thenAnswer((Answer<ResultSet>) invocation -> { + loads.incrementAndGet(); + release.await(); + ResultSet rs = mock(ResultSet.class); + when(rs.next()).thenReturn(true, false); + when(rs.getBytes(1)).thenReturn(PVectorFloat.INSTANCE.toBytes(new float[] { 3, 4 })); + return rs; + }); + int callers = 8; + CachedCentroids[] results = new CachedCentroids[callers]; + List<Thread> threads = new ArrayList<>(); + for (int i = 0; i < callers; i++) { + int caller = i; + Thread thread = new Thread(() -> { + try { + results[caller] = cache.get(conn, "IDX", 1L, DistanceMetric.L2); + } catch (SQLException e) { + throw new RuntimeException(e); + } + }); + threads.add(thread); + thread.start(); + } + // Wait until each caller blocks, in the load or on the load that a different caller started + for (Thread thread : threads) { + while ( + thread.getState() != Thread.State.WAITING && thread.getState() != Thread.State.TIMED_WAITING + ) { + assertTrue("Caller finished before the load was released", thread.isAlive()); + Thread.yield(); + } + } + release.countDown(); + for (Thread thread : threads) { + thread.join(); + } + assertEquals(1, loads.get()); + verify(ps, times(1)).executeQuery(); + for (CachedCentroids result : results) { + assertSame(results[0], result); + } + } + + @Test + public void testInvalidateDropsEveryGenerationOfOneIndex() { + VectorCentroidCache cache = new VectorCentroidCache(MAX_BYTES); + CachedCentroids c = new CachedCentroids(centroids(new float[] { 1 }), DistanceMetric.L2); + cache.put("A", 1L, c); + cache.put("A", 2L, c); + cache.put("B", 1L, c); + cache.invalidate("A"); + assertNull(cache.getIfPresent("A", 1L)); + assertNull(cache.getIfPresent("A", 2L)); + assertSame(c, cache.getIfPresent("B", 1L)); + } + + /** + * Verifies that the cache bound applies to the bytes of the centroid vectors, not to the number + * of models. A model that uses most of the bound stays in the cache. When the total exceeds the + * bound, the cache evicts the least recently used model. + */ + @Test + public void testSizeBoundedByBytes() { + // A two-dimensional centroid uses 8 bytes, so 10 centroids use 80 bytes and 2 use 16 bytes + CachedCentroids large = + new CachedCentroids(Collections.nCopies(10, new float[] { 1, 2 }), DistanceMetric.L2); + CachedCentroids small = + new CachedCentroids(Collections.nCopies(2, new float[] { 1, 2 }), DistanceMetric.L2); + VectorCentroidCache cache = new VectorCentroidCache(100); + cache.put("A", 1L, large); + assertSame(large, cache.getIfPresent("A", 1L)); + cache.put("B", 1L, small); + assertSame(small, cache.getIfPresent("B", 1L)); + assertSame(large, cache.getIfPresent("A", 1L)); + // The total of 112 bytes exceeds the bound, so the cache evicts the least recently used model + cache.put("C", 1L, small); + assertNull(cache.getIfPresent("B", 1L)); + assertSame(large, cache.getIfPresent("A", 1L)); + assertSame(small, cache.getIfPresent("C", 1L)); + cache.put("D", 1L, large); + assertNull(cache.getIfPresent("A", 1L)); + assertEquals(2, cache.size()); + } +} diff --git a/phoenix-core/src/test/java/org/apache/phoenix/index/vector/CentroidManagerTest.java b/phoenix-core/src/test/java/org/apache/phoenix/index/vector/CentroidManagerTest.java new file mode 100644 index 0000000000..4ef8c4508a --- /dev/null +++ b/phoenix-core/src/test/java/org/apache/phoenix/index/vector/CentroidManagerTest.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.index.vector; + +import static org.junit.Assert.assertEquals; + +import org.apache.phoenix.util.EnvironmentEdge; +import org.apache.phoenix.util.EnvironmentEdgeManager; +import org.junit.After; +import org.junit.Test; + +public class CentroidManagerTest { + + @After + public void tearDown() { + EnvironmentEdgeManager.reset(); + } + + private static void setClock(final long now) { + EnvironmentEdgeManager.injectEdge(new EnvironmentEdge() { + @Override + public long currentTime() { + return now; + } + }); + } + + /** + * Verifies that the first generation ID is the wall clock time. A recreated index then gets + * generation IDs that differ from the IDs of the index it replaces. + */ + @Test + public void testFirstGenerationIsCurrentTime() { + setClock(1_700_000_000_000L); + assertEquals(1_700_000_000_000L, CentroidManager.nextGeneration(null)); + } + + /** + * Verifies that the next generation ID is always larger than the current ID, also when the clock + * is not ahead of the current ID. + */ + @Test + public void testNextGenerationStrictlyIncreases() { + setClock(1_000L); + assertEquals(5_000L, CentroidManager.nextGeneration(4_999L)); + assertEquals(1_001L, CentroidManager.nextGeneration(1_000L)); + setClock(2_000L); + assertEquals(2_000L, CentroidManager.nextGeneration(1_000L)); + } +} diff --git a/phoenix-core/src/test/java/org/apache/phoenix/index/vector/KMeansTrainerTest.java b/phoenix-core/src/test/java/org/apache/phoenix/index/vector/KMeansTrainerTest.java new file mode 100644 index 0000000000..fed4d71808 --- /dev/null +++ b/phoenix-core/src/test/java/org/apache/phoenix/index/vector/KMeansTrainerTest.java @@ -0,0 +1,585 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.index.vector; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Random; +import org.apache.phoenix.expression.function.VectorDistanceUtil; +import org.apache.phoenix.optimize.DistanceMetric; +import org.junit.Test; + +public class KMeansTrainerTest { + + private static final double EPSILON = 1e-5; + + /** Returns 2D vectors with Gaussian noise around each given center. */ + private static List<float[]> generate2DGaussianClusters(float[][] centers, int pointsPerCluster, + double stdDev, Random rng) { + List<float[]> dataset = new ArrayList<>(centers.length * pointsPerCluster); + for (float[] center : centers) { + for (int i = 0; i < pointsPerCluster; i++) { + float x = (float) (center[0] + rng.nextGaussian() * stdDev); + float y = (float) (center[1] + rng.nextGaussian() * stdDev); + dataset.add(new float[] { x, y }); + } + } + return dataset; + } + + /** Verifies that training finds the centers of well-separated Gaussian clusters. */ + @Test + public void testClusterRecovery() { + float[][] trueCenters = + new float[][] { { 10.0f, 0.0f }, { 0.0f, 10.0f }, { -10.0f, 0.0f }, { 0.0f, -10.0f } }; + Random rng = new Random(42); + List<float[]> dataset = generate2DGaussianClusters(trueCenters, 100, 0.5, rng); + + KMeansConfig config = KMeansConfig.newBuilder().randomSeed(42L).build(); + KMeansResult result = KMeansTrainer.train(dataset, 4, config); + + assertEquals(4, result.getEffectiveK()); + boolean[] matched = new boolean[4]; + for (float[] centroid : result.getCentroids()) { + int matchIdx = -1; + for (int c = 0; c < 4; c++) { + double dist = + Math.sqrt(VectorDistanceUtil.scalarL2DistanceSquared(centroid, trueCenters[c], 2)); + if (dist < 2.0 && !matched[c]) { + matchIdx = c; + break; + } + } + assertTrue("Centroid " + Arrays.toString(centroid) + " did not match any remaining center", + matchIdx >= 0); + matched[matchIdx] = true; + } + for (int c = 0; c < 4; c++) { + assertTrue("Center " + c + " was not matched", matched[c]); + } + } + + /** + * Verifies that training with k-means++ seeds converges within 50 iterations for 10 different + * random seeds. + */ + @Test + public void testKMeansPlusPlusInitializationQuality() { + float[][] trueCenters = + new float[][] { { 10.0f, 0.0f }, { 0.0f, 10.0f }, { -10.0f, 0.0f }, { 0.0f, -10.0f } }; + Random rng = new Random(100); + List<float[]> dataset = generate2DGaussianClusters(trueCenters, 100, 0.5, rng); + + for (int run = 0; run < 10; run++) { + KMeansConfig config = + KMeansConfig.newBuilder().randomSeed(1000L + run * 17).maxIterations(50).build(); + KMeansResult result = KMeansTrainer.train(dataset, 4, config); + assertTrue("Run " + run + " did not converge in <= 50 iterations (took " + + result.getIterations() + ")", result.isConverged()); + assertTrue(result.getIterations() <= 50); + } + } + + /** Verifies that a single cluster converges to the mean of the full dataset. */ + @Test + public void testSingleClusterDegenerateCase() { + Random rng = new Random(42); + List<float[]> dataset = new ArrayList<>(200); + double sumX = 0.0; + double sumY = 0.0; + for (int i = 0; i < 200; i++) { + float x = (float) (rng.nextGaussian() * 3.0 + 7.0); + float y = (float) (rng.nextGaussian() * 3.0 - 5.0); + sumX += x; + sumY += y; + dataset.add(new float[] { x, y }); + } + + KMeansResult result = KMeansTrainer.train(dataset, 1); + assertEquals(1, result.getEffectiveK()); + assertEquals(1, result.getRequestedK()); + float[] centroid = result.getCentroids().get(0); + assertEquals(sumX / 200.0, centroid[0], 0.01); + assertEquals(sumY / 200.0, centroid[1], 0.01); + } + + /** + * Verifies that the skew metrics show a skewed clustering, and that the split heuristic after + * training decreases the coefficient of variation of the cluster sizes. + */ + @Test + public void testSkewMetricComputationAndSplitHeuristic() { + Random rng = new Random(42); + List<float[]> dataset = new ArrayList<>(1000); + // One large cluster and three small clusters make the dataset skewed + for (int i = 0; i < 900; i++) { + dataset.add(new float[] { (float) (20.0 + rng.nextGaussian() * 0.5), + (float) (20.0 + rng.nextGaussian() * 0.5) }); + } + for (int i = 0; i < 33; i++) { + dataset.add(new float[] { (float) (-20.0 + rng.nextGaussian() * 0.5), + (float) (20.0 + rng.nextGaussian() * 0.5) }); + } + for (int i = 0; i < 33; i++) { + dataset.add(new float[] { (float) (-20.0 + rng.nextGaussian() * 0.5), + (float) (-20.0 + rng.nextGaussian() * 0.5) }); + } + for (int i = 0; i < 34; i++) { + dataset.add(new float[] { (float) (20.0 + rng.nextGaussian() * 0.5), + (float) (-20.0 + rng.nextGaussian() * 0.5) }); + } + + KMeansConfig config = KMeansConfig.newBuilder().randomSeed(42L).enableSplitHeuristic(true) + .splitThresholdMultiplier(2.0).build(); + + KMeansResult result = KMeansTrainer.train(dataset, 4, config); + + assertNotNull(result.getPreSplitSkewMetrics()); + double preSplitCv = result.getPreSplitSkewMetrics().getCoefficientOfVariation(); + assertTrue("Pre-split CV must be > 0.5, got " + preSplitCv, preSplitCv > 0.5); + + assertTrue("Split heuristic should have triggered", result.hasSplit()); + assertTrue("Effective k must be >= requested k", + result.getEffectiveK() >= result.getRequestedK()); + + assertNotNull(result.getPostSplitSkewMetrics()); + double postSplitCv = result.getPostSplitSkewMetrics().getCoefficientOfVariation(); + assertTrue( + "Post-split CV (" + postSplitCv + ") should be lower than pre-split CV (" + preSplitCv + ")", + postSplitCv < preSplitCv); + } + + /** + * Verifies that no cluster stays empty after the split of a cluster of duplicate points. Both + * halves of the split are at the same distance from each duplicate, so the tie can leave one half + * empty. The rebalance after the split must reseed that half. + */ + @Test + public void testSplitRebalanceReseedsPhantomEmptyCentroid() { + List<float[]> dataset = new ArrayList<>(210); + for (int i = 0; i < 200; i++) { + dataset.add(new float[] { 5.0f, 5.0f }); + } + for (int i = 0; i < 10; i++) { + dataset.add(new float[] { -5.0f, -5.0f }); + } + + KMeansConfig config = KMeansConfig.newBuilder().randomSeed(42L).enableSplitHeuristic(true) + .splitThresholdMultiplier(1.1).splitRebalanceIterations(1).build(); + + KMeansResult result = KMeansTrainer.train(dataset, 2, config); + + assertTrue("Split heuristic should have triggered", result.hasSplit()); + assertEquals(3, result.getEffectiveK()); + + int[] clusterSizes = result.getClusterSizes(); + int total = 0; + for (int size : clusterSizes) { + assertTrue( + "No split cluster should be permanently empty, got sizes " + Arrays.toString(clusterSizes), + size > 0); + total += size; + } + assertEquals(210, total); + } + + /** + * Verifies that COSINE training uses spherical k-means and gives unit-length centroids that match + * the cluster directions. + */ + @Test + public void testCosineMetricTraining() { + float[][] directions = new float[][] { { 1.0f, 0.0f, 0.0f }, { 0.0f, 1.0f, 0.0f }, + { 0.0f, 0.0f, 1.0f }, { (float) (-1.0 / Math.sqrt(3.0)), (float) (-1.0 / Math.sqrt(3.0)), + (float) (-1.0 / Math.sqrt(3.0)) } }; + + Random rng = new Random(42); + List<float[]> dataset = new ArrayList<>(400); + for (float[] dir : directions) { + for (int i = 0; i < 100; i++) { + float[] v = new float[] { (float) (dir[0] + (rng.nextFloat() - 0.5f) * 0.05f), + (float) (dir[1] + (rng.nextFloat() - 0.5f) * 0.05f), + (float) (dir[2] + (rng.nextFloat() - 0.5f) * 0.05f) }; + dataset.add(KMeansTrainer.l2Normalize(v)); + } + } + + KMeansConfig config = + KMeansConfig.newBuilder().distanceMetric(DistanceMetric.COSINE).randomSeed(42L).build(); + + KMeansResult result = KMeansTrainer.train(dataset, 4, config); + assertEquals(4, result.getEffectiveK()); + + boolean[] matched = new boolean[4]; + for (float[] centroid : result.getCentroids()) { + double normSq = 0.0; + for (float val : centroid) { + normSq += (double) val * val; + } + double norm = Math.sqrt(normSq); + assertEquals("Centroid must be unit-normalized", 1.0, norm, 1e-5); + + int matchIdx = -1; + for (int d = 0; d < 4; d++) { + double cosDist = VectorDistanceUtil.scalarCosineDistance(centroid, directions[d], 3); + if (cosDist < 0.05 && !matched[d]) { + matchIdx = d; + break; + } + } + assertTrue("Centroid did not match any directional cluster with cosine dist < 0.05", + matchIdx >= 0); + matched[matchIdx] = true; + } + for (int d = 0; d < 4; d++) { + assertTrue("Direction " + d + " was not recovered", matched[d]); + } + } + + /** + * Verifies that INNER_PRODUCT training uses the arithmetic mean as the centroid and does not + * normalize it. + */ + @Test + public void testInnerProductMetricTraining() { + Random rng = new Random(42); + List<float[]> dataset = new ArrayList<>(200); + for (int i = 0; i < 100; i++) { + dataset.add(new float[] { (float) (10.0 + rng.nextGaussian() * 0.5), + (float) (10.0 + rng.nextGaussian() * 0.5) }); + } + for (int i = 0; i < 100; i++) { + dataset.add(new float[] { (float) (-10.0 + rng.nextGaussian() * 0.5), + (float) (-10.0 + rng.nextGaussian() * 0.5) }); + } + + KMeansConfig config = KMeansConfig.newBuilder().distanceMetric(DistanceMetric.INNER_PRODUCT) + .randomSeed(42L).build(); + + KMeansResult result = KMeansTrainer.train(dataset, 2, config); + assertEquals(2, result.getEffectiveK()); + + for (float[] centroid : result.getCentroids()) { + double norm = + Math.sqrt((double) centroid[0] * centroid[0] + (double) centroid[1] * centroid[1]); + // The arithmetic mean keeps the magnitude of the vectors + assertTrue("Centroids must not be unit-normalized for INNER_PRODUCT", norm > 5.0); + } + + boolean hasPositive = false; + boolean hasNegative = false; + for (float[] centroid : result.getCentroids()) { + if (centroid[0] > 5.0 && centroid[1] > 5.0) { + hasPositive = true; + } + if (centroid[0] < -5.0 && centroid[1] < -5.0) { + hasNegative = true; + } + } + assertTrue("Positive cluster center not recovered", hasPositive); + assertTrue("Negative cluster center not recovered", hasNegative); + } + + /** + * Verifies the assignment distance for each metric. INNER_PRODUCT uses L2 distance, so a vector + * goes to the nearest centroid and not to the centroid with the largest inner product. + */ + @Test + public void testAssignmentDistanceByMetric() { + float[] v = new float[] { 10f, 0f }; + float[] near = new float[] { 9f, 0f }; + float[] large = new float[] { 100f, 0f }; + assertTrue(KMeansTrainer.assignmentDistance(DistanceMetric.INNER_PRODUCT, v, near) + < KMeansTrainer.assignmentDistance(DistanceMetric.INNER_PRODUCT, v, large)); + assertEquals(1.0, KMeansTrainer.assignmentDistance(DistanceMetric.L2, v, near), EPSILON); + assertEquals(VectorDistanceUtil.scalarCosineDistance(v, new float[] { 0f, 1f }, 2), + KMeansTrainer.assignmentDistance(DistanceMetric.COSINE, v, new float[] { 0f, 1f }), EPSILON); + } + + /** + * Verifies that training converges by the relative change in distortion, and that distortion does + * not increase from one iteration to the next. + */ + @Test + public void testConvergenceByRelativeDistortion() { + float[][] trueCenters = + new float[][] { { 10.0f, 0.0f }, { 0.0f, 10.0f }, { -10.0f, 0.0f }, { 0.0f, -10.0f } }; + Random rng = new Random(42); + List<float[]> dataset = generate2DGaussianClusters(trueCenters, 100, 0.5, rng); + + KMeansConfig config = + KMeansConfig.newBuilder().randomSeed(42L).convergenceThreshold(1e-4).build(); + + KMeansResult result = KMeansTrainer.train(dataset, 4, config); + assertTrue(result.isConverged()); + assertTrue(result.getFinalDistortion() > 0.0); + assertFalse(Double.isInfinite(result.getFinalDistortion())); + assertFalse(Double.isNaN(result.getFinalDistortion())); + + List<Double> history = result.getDistortionHistory(); + assertTrue(history.size() >= 2); + for (int i = 1; i < history.size(); i++) { + assertTrue("Distortion should not increase at iteration " + i + ": prev=" + history.get(i - 1) + + ", curr=" + history.get(i), history.get(i) <= history.get(i - 1) + EPSILON); + } + } + + /** + * Verifies that training rejects a null or empty vector list, a cluster count out of range, and + * vectors of different dimensions. + */ + @Test + public void testValidationAndEdgeCases() { + List<float[]> validVectors = Arrays.asList(new float[] { 1f, 2f }, new float[] { 3f, 4f }); + + try { + KMeansTrainer.train((List<float[]>) null, 2); + fail("Expected IllegalArgumentException on null vector list"); + } catch (IllegalArgumentException expected) { + } + + try { + KMeansTrainer.train(Collections.<float[]> emptyList(), 2); + fail("Expected IllegalArgumentException on empty vector list"); + } catch (IllegalArgumentException expected) { + } + + try { + KMeansTrainer.train(validVectors, 0); + fail("Expected IllegalArgumentException on k <= 0"); + } catch (IllegalArgumentException expected) { + } + + try { + KMeansTrainer.train(validVectors, -1); + fail("Expected IllegalArgumentException on k < 0"); + } catch (IllegalArgumentException expected) { + } + + try { + KMeansTrainer.train(validVectors, 3); + fail("Expected IllegalArgumentException on k > N"); + } catch (IllegalArgumentException expected) { + } + + List<float[]> mismatched = Arrays.asList(new float[] { 1f, 2f }, new float[] { 3f, 4f, 5f }); + try { + KMeansTrainer.train(mismatched, 2); + fail("Expected IllegalArgumentException on mismatched dimensions"); + } catch (IllegalArgumentException expected) { + } + } + + /** Verifies that each centroid is one of the input points when k equals the sample size. */ + @Test + public void testKEqualsN() { + List<float[]> dataset = new ArrayList<>(10); + for (int i = 0; i < 10; i++) { + dataset.add(new float[] { (float) (i * 2.0), (float) (i * 3.0) }); + } + + KMeansConfig config = KMeansConfig.newBuilder().randomSeed(42L).build(); + KMeansResult result = KMeansTrainer.train(dataset, 10, config); + + assertEquals(10, result.getEffectiveK()); + boolean[] matched = new boolean[10]; + for (float[] c : result.getCentroids()) { + int matchIdx = -1; + for (int i = 0; i < 10; i++) { + float[] orig = dataset.get(i); + if (Math.abs(c[0] - orig[0]) < 1e-4 && Math.abs(c[1] - orig[1]) < 1e-4 && !matched[i]) { + matchIdx = i; + break; + } + } + assertTrue("Centroid must match an original point: " + Arrays.toString(c), matchIdx >= 0); + matched[matchIdx] = true; + } + } + + /** Verifies that ClusterSkewMetrics serializes to bytes and back without change. */ + @Test + public void testClusterSkewMetricsSerialization() { + int[] sizes = new int[] { 5, 10, 20, 50, 100, 250, 500 }; + ClusterSkewMetrics m1 = ClusterSkewMetrics.compute(sizes); + + byte[] bytes = m1.toBytes(); + assertNotNull(bytes); + assertTrue(bytes.length > 0); + + ClusterSkewMetrics m2 = ClusterSkewMetrics.fromBytes(bytes); + assertEquals(m1, m2); + assertEquals(m1.getMin(), m2.getMin()); + assertEquals(m1.getMax(), m2.getMax()); + assertEquals(m1.getMean(), m2.getMean(), EPSILON); + assertEquals(m1.getP95(), m2.getP95(), EPSILON); + assertEquals(m1.getStdDev(), m2.getStdDev(), EPSILON); + assertEquals(m1.getCoefficientOfVariation(), m2.getCoefficientOfVariation(), EPSILON); + assertEquals(m1.getTotalVectors(), m2.getTotalVectors()); + assertEquals(m1.getK(), m2.getK()); + assertArrayEquals(m1.getClusterSizes(), m2.getClusterSizes()); + } + + /** + * Verifies that reservoirSample returns the correct sample size and selects each element with + * equal probability. + */ + @Test + public void testReservoirSampling() { + List<float[]> streamSmall = new ArrayList<>(5); + for (int i = 0; i < 5; i++) { + streamSmall.add(new float[] { (float) i }); + } + List<float[]> sample1 = + KMeansTrainer.reservoirSample(streamSmall.iterator(), 10, new Random(42)); + assertEquals(5, sample1.size()); + + List<float[]> streamLarge = new ArrayList<>(20); + for (int i = 0; i < 20; i++) { + streamLarge.add(new float[] { (float) i }); + } + List<float[]> sample2 = + KMeansTrainer.reservoirSample(streamLarge.iterator(), 10, new Random(42)); + assertEquals(10, sample2.size()); + + int n = 10; + int k = 3; + int runs = 10000; + int[] counts = new int[n]; + Random rng = new Random(12345); + + for (int r = 0; r < runs; r++) { + List<float[]> items = new ArrayList<>(n); + for (int i = 0; i < n; i++) { + items.add(new float[] { (float) i }); + } + List<float[]> sample = KMeansTrainer.reservoirSample(items.iterator(), k, rng); + for (float[] v : sample) { + counts[(int) v[0]]++; + } + } + + // Each count has a binomial distribution with mean 3000 and a standard deviation of about 46. + // The bounds are approximately 6.5 standard deviations from the mean. + for (int i = 0; i < n; i++) { + assertTrue("Element " + i + " count (" + counts[i] + ") outside expected range [2700, 3300]", + counts[i] >= 2700 && counts[i] <= 3300); + } + } + + /** Verifies that reservoirSample ignores null elements. */ + @Test + public void testReservoirSamplingSkipsNulls() { + List<float[]> stream = new ArrayList<>(20); + int expectedNulls = 7; + for (int i = 0; i < 20; i++) { + if (i % 3 == 0 && expectedNulls > 0) { + stream.add(null); + expectedNulls--; + } else { + stream.add(new float[] { (float) i, (float) (i * 2) }); + } + } + + List<float[]> sample = KMeansTrainer.reservoirSample(stream.iterator(), 10, new Random(42)); + assertEquals(10, sample.size()); + for (float[] v : sample) { + assertNotNull("Sample must not contain null vectors", v); + } + } + + /** Verifies that two training runs with the same random seed give the same centroids. */ + @Test + public void testDeterministicReproducibility() { + float[][] trueCenters = + new float[][] { { 10.0f, 0.0f }, { 0.0f, 10.0f }, { -10.0f, 0.0f }, { 0.0f, -10.0f } }; + Random rng = new Random(42); + List<float[]> dataset = generate2DGaussianClusters(trueCenters, 100, 0.5, rng); + + KMeansConfig config = KMeansConfig.newBuilder().randomSeed(99999L).build(); + + KMeansResult r1 = KMeansTrainer.train(dataset, 4, config); + KMeansResult r2 = KMeansTrainer.train(dataset, 4, config); + + assertEquals(r1.getEffectiveK(), r2.getEffectiveK()); + assertEquals(r1.getIterations(), r2.getIterations()); + for (int c = 0; c < r1.getEffectiveK(); c++) { + assertArrayEquals("Centroid " + c + " must match across runs with same seed", + r1.getCentroids().get(c), r2.getCentroids().get(c), 0.0f); + } + } + + /** Verifies that training finds each of the four cluster centers in 128 dimensions. */ + @Test(timeout = 10000) + public void testHighDimensionalTraining() { + int dim = 128; + int k = 4; + int pointsPerCluster = 125; + Random rng = new Random(42); + + float[][] centers = new float[k][dim]; + for (int c = 0; c < k; c++) { + for (int d = 0; d < dim; d++) { + centers[c][d] = (c == (d % k)) ? 20.0f : 0.0f; + } + } + + List<float[]> dataset = new ArrayList<>(k * pointsPerCluster); + for (int c = 0; c < k; c++) { + for (int i = 0; i < pointsPerCluster; i++) { + float[] v = new float[dim]; + for (int d = 0; d < dim; d++) { + v[d] = (float) (centers[c][d] + rng.nextGaussian() * 0.5); + } + dataset.add(v); + } + } + + KMeansConfig config = KMeansConfig.newBuilder().randomSeed(42L).maxIterations(50).build(); + KMeansResult result = KMeansTrainer.train(dataset, k, config); + + assertEquals(k, result.getEffectiveK()); + + boolean[] matched = new boolean[k]; + for (float[] centroid : result.getCentroids()) { + int matchIdx = -1; + for (int c = 0; c < k; c++) { + double dist = + Math.sqrt(VectorDistanceUtil.scalarL2DistanceSquared(centroid, centers[c], dim)); + if (dist < 5.0 && !matched[c]) { + matchIdx = c; + break; + } + } + assertTrue("Centroid did not match any true high-dimensional center", matchIdx >= 0); + matched[matchIdx] = true; + } + for (int c = 0; c < k; c++) { + assertTrue("Center " + c + " was not matched in high-dimensional test", matched[c]); + } + } +} diff --git a/phoenix-core/src/test/java/org/apache/phoenix/index/vector/VectorIndexTrainerTest.java b/phoenix-core/src/test/java/org/apache/phoenix/index/vector/VectorIndexTrainerTest.java new file mode 100644 index 0000000000..a35849a7ab --- /dev/null +++ b/phoenix-core/src/test/java/org/apache/phoenix/index/vector/VectorIndexTrainerTest.java @@ -0,0 +1,109 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.phoenix.index.vector; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.phoenix.jdbc.PhoenixConnection; +import org.apache.phoenix.schema.PColumn; +import org.apache.phoenix.schema.PNameFactory; +import org.apache.phoenix.schema.PTable; +import org.apache.phoenix.schema.types.PVectorDouble; +import org.junit.Test; + +public class VectorIndexTrainerTest { + + /** + * Verifies that the training sample does not include a VECTOR(DOUBLE) value with an element + * beyond float range. That element narrows to infinity and can corrupt the trained centroids. + */ + @Test + public void testOutOfFloatRangeVectorLeftOutOfSample() throws SQLException { + KMeansResult result = VectorIndexTrainer + .train(connection(new double[] { 0, 0 }, new double[] { 0, 1 }, new double[] { 1e300, 0 }, + new double[] { 10, 10 }, new double[] { 10, 11 }), dataTable(), index(2)); + for (float[] centroid : result.getCentroids()) { + for (float x : centroid) { + assertTrue(Float.isFinite(x)); + } + } + assertEquals(4, result.getAssignments().length); + } + + /** + * Verifies that train() returns null, and so defers training, when the sample has fewer trainable + * vectors than the index has lists. The row count is sufficient here, but only one vector is + * finite in float range. + */ + @Test + public void testTooFewTrainableVectorsDefersTraining() throws SQLException { + assertNull(VectorIndexTrainer.train( + connection(new double[] { 1e300, 0 }, new double[] { 0, -1e300 }, new double[] { 1, 1 }), + dataTable(), index(2))); + } + + private static PhoenixConnection connection(Object... vectors) throws SQLException { + Statement stmt = mock(Statement.class); + when(stmt.executeQuery(anyString())) + .thenAnswer(inv -> ((String) inv.getArgument(0)).startsWith("SELECT COUNT(*)") + ? rows((long) vectors.length) + : rows(vectors)); + PhoenixConnection conn = mock(PhoenixConnection.class); + when(conn.createStatement()).thenReturn(stmt); + return conn; + } + + private static ResultSet rows(Object... values) throws SQLException { + ResultSet rs = mock(ResultSet.class); + AtomicInteger row = new AtomicInteger(-1); + when(rs.next()).thenAnswer(inv -> row.incrementAndGet() < values.length); + when(rs.getObject(1)).thenAnswer(inv -> values[row.get()]); + when(rs.getLong(1)).thenAnswer(inv -> (Long) values[row.get()]); + return rs; + } + + private static PTable dataTable() { + PTable table = mock(PTable.class); + when(table.getName()).thenReturn(PNameFactory.newName("T")); + return table; + } + + private static PTable index(int lists) { + PColumn column = mock(PColumn.class); + when(column.getFamilyName()).thenReturn(PNameFactory.newName("0")); + when(column.getDataType()).thenReturn(PVectorDouble.INSTANCE); + when(column.getExpressionStr()).thenReturn("V"); + PTable index = mock(PTable.class); + when(index.getName()).thenReturn(PNameFactory.newName("IDX")); + when(index.getColumns()).thenReturn(Collections.singletonList(column)); + when(index.getVectorIvfLists()).thenReturn(lists); + when(index.getVectorIvfSampleSize()).thenReturn(100); + when(index.getVectorDistanceMetric()).thenReturn("L2"); + return index; + } +} diff --git a/phoenix-core/src/test/java/org/apache/phoenix/optimize/VectorSearchUtilTest.java b/phoenix-core/src/test/java/org/apache/phoenix/optimize/VectorSearchUtilTest.java index a572e7279d..ce880a49c2 100644 --- a/phoenix-core/src/test/java/org/apache/phoenix/optimize/VectorSearchUtilTest.java +++ b/phoenix-core/src/test/java/org/apache/phoenix/optimize/VectorSearchUtilTest.java @@ -261,6 +261,32 @@ public class VectorSearchUtilTest extends BaseConnectionlessQueryTest { } } + @Test + public void testDistanceRewrittenToIndexColumnKeepsNullsLast() throws Exception { + try (PhoenixConnection conn = (PhoenixConnection) DriverManager.getConnection(getUrl())) { + String q = "ARRAY[1.0, 0.0, 0.0]"; + conn.createStatement() + .execute("CREATE TABLE t_fidx (pk INTEGER PRIMARY KEY, v VECTOR(FLOAT, 3))"); + conn.createStatement().execute("CREATE INDEX i_fidx ON t_fidx (L2_DISTANCE(v, " + q + "))"); + // The functional index rewrites the distance to an index column, whether the query orders + // by the distance, by an alias, or by an ordinal. The index plan must still sort rows that + // have no vector last, as the data plan does. The row key order of the index puts them first. + for (String hint : new String[] { "/*+ INDEX(t_fidx i_fidx) */", "/*+ NO_INDEX */" }) { + for (String sql : new String[] { + "SELECT " + hint + " pk FROM t_fidx ORDER BY L2_DISTANCE(v, " + q + ") LIMIT 5", + "SELECT " + hint + " pk, L2_DISTANCE(v, " + q + ") d FROM t_fidx ORDER BY d LIMIT 5", + "SELECT " + hint + " pk, L2_DISTANCE(v, " + q + ") FROM t_fidx ORDER BY 2 LIMIT 5" }) { + org.apache.phoenix.compile.QueryPlan plan = conn.prepareStatement(sql) + .unwrap(org.apache.phoenix.jdbc.PhoenixPreparedStatement.class).optimizeQuery(); + assertEquals(sql, hint.contains("NO_INDEX") ? "T_FIDX" : "I_FIDX", + plan.getTableRef().getTable().getTableName().getString()); + assertEquals(sql, 1, plan.getOrderBy().getOrderByExpressions().size()); + assertTrue(sql, plan.getOrderBy().getOrderByExpressions().get(0).isNullsLast()); + } + } + } + } + @Test public void testLimitAsBindVariable() throws Exception { String sql = "SELECT * FROM t ORDER BY L2_DISTANCE(v, ?) ASC LIMIT ?";
