diff --git a/hudi-common/src/main/avro/HoodieMetadata.avsc b/hudi-common/src/main/avro/HoodieMetadata.avsc index 84dc97dc67206..880c17bd4f9f2 100644 --- a/hudi-common/src/main/avro/HoodieMetadata.avsc +++ b/hudi-common/src/main/avro/HoodieMetadata.avsc @@ -554,6 +554,147 @@ } ], "default" : null + }, + { + "name": "VectorIndexMetadata", + "doc": "Typed metadata records for the MDT-backed IVF + RaBitQ vector index.", + "type": [ + "null", + { + "type": "record", + "name": "HoodieVectorIndexActiveManifest", + "fields": [ + {"name": "indexVersion", "type": "int", "doc": "Persisted vector-index format version."}, + {"name": "activeGeneration", "type": ["null", "int"], "default": null, + "doc": "Reader-visible generation ordinal, or null before first activation."} + ] + }, + { + "type": "record", + "name": "HoodieVectorIndexManifest", + "fields": [ + {"name": "indexVersion", "type": "int"}, + {"name": "generationId", "type": "string", "doc": "Creating Hudi instant + optional human tag. The gen ordinal lives in the key."}, + {"name": "state", "type": "string", "doc": "BUILDING | ACTIVE | RETIRED. Readers serve only the generation referenced by the singleton active manifest."}, + {"name": "dim", "type": "int"}, + {"name": "dimPadded", "type": "int", "doc": "nextPow2(dim) for hadamard quantizers, else dim. Authoritative for codeRowBytes."}, + {"name": "codeRowBytes", "type": "int", "doc": "ceil(dimPadded/64)*8. Long-aligned per-plane row width."}, + {"name": "bitsTotal", "type": "int", "doc": "B = 1 + numExPlanes."}, + {"name": "numExPlanes", "type": "int"}, + {"name": "numClusters", "type": "int"}, + {"name": "shardCount", "type": "int", "doc": "Initial per-cluster posting shard geometry frozen for this generation."}, + {"name": "fileGroupCount", "type": "int", "doc": "MDT file-group geometry frozen for this generation; readers and writers never re-derive it."}, + {"name": "metric", "type": "string", "doc": "L2 | DOT | COSINE."}, + {"name": "assumeNormalized", "type": "boolean"}, + {"name": "residualEncoding", "type": "boolean", "default": false, "doc": "True when posting codes encode x - centroid and readers must use residual scoring."}, + {"name": "vectorColumn", "type": "string", "doc": "Authoritative base-table vector column used for bootstrap and exact rerank."}, + {"name": "targetBlockBytes", "type": "int", "doc": "Config input, e.g. 524288."}, + {"name": "vectorsPerBlock", "type": "int", "doc": "Resolved N. Frozen per generation; readers never re-derive."}, + {"name": "blockFormatVersion", "type": "int", "doc": "Expected posting-block layout version for this generation."}, + {"name": "factorVersion", "type": "int", "doc": "RaBitQ scalar-factor semantics version. Readers reject unsupported values."}, + {"name": "kappa", "type": "double", "doc": "Generation-scoped pass-1 error scale."}, + {"name": "gMin", "type": "double", "doc": "Generation-scoped normalized alignment floor."}, + {"name": "eps1Max", "type": "double", "doc": "Maximum permitted relative pass-1 error."}, + {"name": "epsNRel", "type": "double", "doc": "Relative residual-norm floor."}, + {"name": "centroidChunkCount", "type": "int", "doc": "Expected number of size-bounded centroid chunks."}, + {"name": "centroidChecksum", "type": ["null", "string"], "default": null, + "doc": "Checksum of the canonical centroid payload used to reject incomplete or mismatched builds."}, + {"name": "splitLimit", "type": "int"}, + {"name": "mergeFloor", "type": "int"}, + {"name": "bootstrapInstant", "type": ["null", "string"], "default": null, + "doc": "Source data-write instant included by the generation bootstrap baseline."}, + {"name": "verifiedFrontier", "type": ["null", "string"], "default": null, + "doc": "Latest source data-write instant with verified contiguous marker coverage from the bootstrap baseline."}, + {"name": "createdTs", "type": "long"} + ] + }, + { + "type": "record", + "name": "HoodieVectorIndexQuantizer", + "fields": [ + {"name": "quantizerType", "type": "string", "doc": "rb-rotation-explicit | rb-rotation-hadamard."}, + {"name": "randomSeed", "type": "long"}, + {"name": "rotationBytes", "type": ["null", "bytes"], "default": null, + "doc": "Explicit path: rows of the D x D float32 LE orthogonal matrix. Null for hadamard."} + ] + }, + { + "type": "record", + "name": "HoodieVectorIndexCentroids", + "fields": [ + {"name": "clusterIds", "type": "bytes", "doc": "k x u32 LE cluster ids in this chunk, parallel to centroidBytes rows."}, + {"name": "centroidBytes", "type": "bytes", "doc": "k x dimPadded float32 LE, row-major, in rotated space."}, + {"name": "clusterRadii", "type": "bytes", "doc": "k x float32 LE max residual norm per cluster."} + ] + }, + { + "type": "record", + "name": "HoodieVectorIndexPostingBlock", + "fields": [ + {"name": "blockFormatVersion", "type": "int"}, + {"name": "numVectors", "type": "int"}, + {"name": "codeRowBytes", "type": "int", "doc": "Duplicated from manifest for self-description; MUST match."}, + {"name": "signPlane", "type": "bytes", "doc": "S1: N x codeRowBytes, row-major. Bit plane B-1 (MSB)."}, + {"name": "exPlanes", "type": "bytes", "doc": "S2: N x numExPlanes x codeRowBytes. Vector-major; planes MSB-first."}, + {"name": "scalarFactors", "type": "bytes", "doc": "S3: fAdd1 | fRescale1 | err1 | fAddEx | fRescaleEx | residualNorm | optional vectorNorm, float32 LE SoA."}, + {"name": "rowLocators", "type": "bytes", "doc": "S4: N x 8 B LE structs: fgDictIdx u16 | instantDictIdx u16 | rowPosition u32."}, + {"name": "fileGroupDict", "type": {"type": "array", "items": "string"}}, + {"name": "instantTimeDict", "type": {"type": "array", "items": "string"}}, + {"name": "partitionDict", "type": {"type": "array", "items": "string"}, "doc": "Parallel indexing with fileGroupDict."}, + {"name": "recordKeyOffsets", "type": "bytes", "doc": "S6a: (N+1) x u32 LE offsets into recordKeyBytes."}, + {"name": "recordKeyBytes", "type": "bytes", "doc": "S6b: concatenated UTF-8 keys. Must remain final field."} + ] + }, + { + "type": "record", + "name": "HoodieVectorIndexPostingDelta", + "fields": [ + {"name": "recordKey", "type": "string"}, + {"name": "binaryCode", "type": "bytes", "doc": "MSB row followed by ex rows, identical row encoding and plane order as blocks."}, + {"name": "fAdd1", "type": "float"}, + {"name": "fRescale1", "type": "float"}, + {"name": "err1", "type": "float"}, + {"name": "fAddEx", "type": "float"}, + {"name": "fRescaleEx", "type": "float"}, + {"name": "residualNorm", "type": "float"}, + {"name": "vectorNorm", "type": ["null", "float"], "default": null}, + {"name": "fileGroupId", "type": "string"}, + {"name": "partitionPath", "type": "string"}, + {"name": "baseInstantTime", "type": "string"}, + {"name": "rowPosition", "type": "long"} + ] + }, + { + "type": "record", + "name": "HoodieVectorIndexClusterStats", + "fields": [ + {"name": "routingVersion", "type": "int", "doc": "Cache invalidation version for atomically published cluster routing changes."}, + {"name": "shardCount", "type": "int"}, + {"name": "fileGroupIds", "type": {"type": "array", "items": "string"}, "doc": "Candidate base-table file groups for this cluster."}, + {"name": "liveCount", "type": "long"}, + {"name": "deltaCount", "type": "long"}, + {"name": "tombstoneCount", "type": "long"}, + {"name": "lastRebalanceInstant", "type": ["null", "string"], "default": null}, + {"name": "lastUpdatedTs", "type": "long"} + ] + }, + { + "type": "record", + "name": "HoodieVectorIndexSourceInstantMarker", + "fields": [ + {"name": "dataInstant", "type": "string", "doc": "Completed source data-write instant incorporated by this generation."} + ] + }, + { + "type": "record", + "name": "HoodieVectorIndexTombstone", + "fields": [ + {"name": "deleteInstant", "type": "string"}, + {"name": "deleteReason", "type": "string", "doc": "SPLIT | MERGE | REBALANCE | GC | MANUAL."} + ] + } + ], + "default" : null } ] } diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/MetricQueryState.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/MetricQueryState.java new file mode 100644 index 0000000000000..cb2586e56aac2 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/MetricQueryState.java @@ -0,0 +1,294 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +import java.io.Serializable; + +/** + * Per-metric query state using residual-query estimation in rotated space (RFC-109 §3). + * + *

The MDT stores metric-neutral residual factors. The padded query is rotated exactly + * once in the constructor to {@code qRot = P(qPad)}. For each probed IVF cluster the scorer + * consumes an already-rotated centroid {@code cRot} and forms the rotated residual query + * {@code wRot = qRot - cRot} by subtraction only — it never rotates {@code q - c} per cluster. + * Because rotation is a linear isometry, {@code P(q - c) = P(q) - P(c)} and + * the inner product {@code <q, c> = <qRot, cRot>}, so exact terms are computed directly + * in rotated space. + */ +public abstract class MetricQueryState implements Serializable { + + private static final long serialVersionUID = 2L; + protected static final double EPS_NORM = 1.0e-9; + + /** Rotated padded query {@code qRot = P(qPad)}; computed once. */ + protected final float[] rotatedQuery; + /** ||q|| (rotation-invariant). */ + protected final double queryNorm; + /** The shared query rotation {@code P}; retained only to rotate centroids once (cached). */ + private final QueryRotation rotation; + + protected MetricQueryState(QueryRotation rotation, float[] paddedQuery) { + this.rotation = rotation; + this.rotatedQuery = validateRotationResult( + rotation.apply(paddedQuery.clone()), paddedQuery.length, "query"); + this.queryNorm = Math.sqrt(sqNorm(rotatedQuery)); + } + + /** + * Apply the shared rotation {@code P} to a raw padded centroid, yielding {@code cRot = P(c)}. + * Callers should invoke this at most once per distinct probed cluster (cached) — never per + * {@code q - c} — and pass the result to {@link #forRotatedCentroid(float[])}. Prefer feeding + * pre-rotated centroids from the index cache (RFC-109 §3) where available. + */ + public final float[] rotateCentroid(float[] paddedCentroid) { + if (paddedCentroid == null) { + throw new IllegalArgumentException("paddedCentroid must not be null"); + } + return validateRotationResult( + rotation.apply(paddedCentroid.clone()), paddedCentroid.length, "centroid"); + } + + /** Whether this metric requires the 7th (VECTOR_NORM) scalar array in blocks. */ + public boolean requiresVectorNorm() { + return false; + } + + /** Reject degenerate queries per the spec's normative rules. */ + public void validate() { + // L2 permits any query (s ~ 0 handled per cluster); overridden where zero-q is invalid. + } + + /** The rotated padded query {@code qRot}; exposed so callers may compute error bounds. */ + public final float[] rotatedQuery() { + return rotatedQuery.clone(); + } + + /** + * Build per-cluster query state from an already-rotated centroid {@code cRot = P(c)}. The float + * kernel consumes {@code wRot = qRot - cRot} and its sum. No rotation is performed here — callers + * must rotate (and preferably cache) centroids once, never per {@code q - c}. + */ + public final ClusterQuery forRotatedCentroid(float[] rotatedCentroid) { + int n = rotatedQuery.length; + if (rotatedCentroid.length != n) { + throw new IllegalArgumentException( + "Rotated centroid dimension mismatch: expected " + n + ", got " + rotatedCentroid.length); + } + float[] residualQuery = new float[n]; + for (int i = 0; i < n; i++) { + residualQuery[i] = rotatedQuery[i] - rotatedCentroid[i]; + } + double wNormSq = sqNorm(residualQuery); + return new ClusterQuery( + residualQuery, + sum(residualQuery), + exactTerms(rotatedCentroid, wNormSq), + wNormSq); + } + + /** Exact per-metric terms; {@code rotatedCentroid} is {@code P(c)}, {@code wNormSq = ||q - c||^2}. */ + protected abstract ExactTerms exactTerms(float[] rotatedCentroid, double wNormSq); + + /** + * Metric-specific ranking value from the residual inner-product estimate {@code ripW}. + * L2 uses squared Euclidean distance to avoid a square root in the posting scan; cosine and + * dot-product use the same units as {@link VectorDistanceMetric}. + */ + public abstract double rankingDistance(double ripW, float centerRip, float residualNorm, + float vectorNorm, ClusterQuery cq); + + /** Optimistic ranking value for pass-1 pruning, using the residual-query norm. */ + public final double optimisticRankingDistance(double ripW, float err1, float centerRip, + float residualNorm, float vectorNorm, + ClusterQuery cq) { + double errAbs = (double) err1 * Math.sqrt(cq.wNormSq); + return rankingDistance(ripW + errAbs, centerRip, residualNorm, vectorNorm, cq); + } + + public static final class ClusterQuery implements Serializable { + private static final long serialVersionUID = 2L; + public final float[] rotatedQuery; + public final float querySum; + public final ExactTerms terms; + public final double wNormSq; + public final boolean queryAtCentroid; + + ClusterQuery(float[] rotatedQuery, float querySum, ExactTerms terms, double wNormSq) { + this.rotatedQuery = rotatedQuery; + this.querySum = querySum; + this.terms = terms; + this.wNormSq = wNormSq; + this.queryAtCentroid = wNormSq < EPS_NORM * EPS_NORM; + } + } + + public static final class ExactTerms implements Serializable { + private static final long serialVersionUID = 2L; + public final double qDotC; + + ExactTerms(double qDotC) { + this.qDotC = qDotC; + } + } + + // -------------------------------------------------------------------------------------- + + public static MetricQueryState create(VectorDistanceMetric metric, QueryRotation rotation, + float[] paddedQuery, + boolean assumeNormalized) { + if (metric == null || rotation == null || paddedQuery == null) { + throw new IllegalArgumentException("metric, rotation, and paddedQuery must not be null"); + } + for (int i = 0; i < paddedQuery.length; i++) { + if (!Float.isFinite(paddedQuery[i])) { + throw new IllegalArgumentException("paddedQuery contains a non-finite value at dimension " + i); + } + } + MetricQueryState state; + switch (metric) { + case L2: + state = new L2QueryState(rotation, paddedQuery); + break; + case DOT_PRODUCT: + state = new DotQueryState(rotation, paddedQuery); + break; + case COSINE: + state = new CosineQueryState(rotation, paddedQuery, assumeNormalized); + break; + default: + throw new IllegalArgumentException("Unsupported metric: " + metric); + } + state.validate(); + return state; + } + + // -------------------------------------------------------------------------------------- + + /** d2 = ||q - c||^2 + n^2 - 2 * ripW. */ + public static final class L2QueryState extends MetricQueryState { + L2QueryState(QueryRotation rotation, float[] paddedQuery) { + super(rotation, paddedQuery); + } + + @Override + protected ExactTerms exactTerms(float[] rotatedCentroid, double wNormSq) { + return new ExactTerms(0.0); + } + + @Override + public double rankingDistance(double ripW, float centerRip, float residualNorm, + float vectorNorm, ClusterQuery cq) { + if (cq.queryAtCentroid) { + return (double) residualNorm * residualNorm; // exact; scanner may skip the kernel + } + return cq.wNormSq + (double) residualNorm * residualNorm - 2.0 * ripW; + } + } + + /** Dot-product distance finish. Zero query rejected. */ + public static final class DotQueryState extends MetricQueryState { + DotQueryState(QueryRotation rotation, float[] paddedQuery) { + super(rotation, paddedQuery); + } + + @Override + public void validate() { + if (queryNorm < 1.0e-9f) { + throw new IllegalArgumentException("DOT_PRODUCT query must be non-zero."); + } + } + + @Override + protected ExactTerms exactTerms(float[] rotatedCentroid, double wNormSq) { + return new ExactTerms(dot(rotatedQuery, rotatedCentroid)); + } + + @Override + public double rankingDistance(double ripW, float centerRip, float residualNorm, + float vectorNorm, ClusterQuery cq) { + return -(cq.terms.qDotC + centerRip + ripW); + } + } + + /** Cosine distance finish. Zero query rejected. */ + public static final class CosineQueryState extends MetricQueryState { + private final boolean assumeNormalized; + + CosineQueryState(QueryRotation rotation, float[] paddedQuery, boolean assumeNormalized) { + super(rotation, paddedQuery); + this.assumeNormalized = assumeNormalized; + } + + @Override + public void validate() { + if (queryNorm < 1.0e-9f) { + throw new IllegalArgumentException("COSINE query must be non-zero."); + } + } + + @Override + public boolean requiresVectorNorm() { + return !assumeNormalized; + } + + @Override + protected ExactTerms exactTerms(float[] rotatedCentroid, double wNormSq) { + return new ExactTerms(dot(rotatedQuery, rotatedCentroid)); + } + + @Override + public double rankingDistance(double ripW, float centerRip, float residualNorm, + float vectorNorm, ClusterQuery cq) { + double xNorm = assumeNormalized ? 1.0 : vectorNorm; + if (xNorm < EPS_NORM) { + return 1.0; // zero vector: orthogonal by convention + } + return 1.0 - (cq.terms.qDotC + centerRip + ripW) / (queryNorm * xNorm); + } + } + + // -------------------------------------------------------------------------------------- + + private static float[] validateRotationResult(float[] result, int expectedLength, String name) { + if (result == null || result.length != expectedLength) { + throw new IllegalArgumentException("Rotation returned an invalid " + name + " dimension"); + } + for (int i = 0; i < result.length; i++) { + if (!Float.isFinite(result[i])) { + throw new IllegalArgumentException("Rotation returned a non-finite " + name + " value at " + i); + } + } + return result.clone(); + } + + private static double sqNorm(float[] v) { + double s = 0.0; + for (float x : v) { + s += (double) x * x; + } + return s; + } + + private static double dot(float[] a, float[] b) { + double s = 0.0; + for (int i = 0; i < a.length; i++) { + s += (double) a[i] * b[i]; + } + return s; + } + + private static float sum(float[] values) { + float sum = 0.0f; + for (float value : values) { + sum += value; + } + return sum; + } + + public interface QueryRotation extends Serializable { + float[] apply(float[] vector); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/PostingBlockBuilder.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/PostingBlockBuilder.java new file mode 100644 index 0000000000000..16d913a2a1834 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/PostingBlockBuilder.java @@ -0,0 +1,346 @@ +/* + * 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.hudi.common.index.vector; + +import org.apache.hudi.avro.model.HoodieVectorIndexPostingBlock; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +import static org.apache.hudi.common.util.ValidationUtils.checkArgument; + +/** + * Streaming-friendly builder for one immutable vector posting block. + */ +public final class PostingBlockBuilder { + + public static final int BLOCK_FORMAT_VERSION = 1; + public static final int SCALAR_FACTOR_COUNT = 6; + public static final int SCALAR_FACTOR_COUNT_WITH_VECTOR_NORM = 7; + public static final int ROW_LOCATOR_BYTES = 8; + + private final int codeRowBytes; + private final int numExPlanes; + private final boolean includeVectorNorm; + private final List rows = new ArrayList<>(); + private final Map fileGroupDict = new LinkedHashMap<>(); + private final Map instantTimeDict = new LinkedHashMap<>(); + private final List partitionDict = new ArrayList<>(); + + public PostingBlockBuilder(int codeRowBytes, int numExPlanes) { + this(codeRowBytes, numExPlanes, false); + } + + public PostingBlockBuilder(int codeRowBytes, int numExPlanes, boolean includeVectorNorm) { + checkArgument(codeRowBytes > 0 && codeRowBytes % Long.BYTES == 0, + "codeRowBytes must be positive and long-aligned: " + codeRowBytes); + checkArgument(numExPlanes >= 0, "numExPlanes must be non-negative: " + numExPlanes); + this.codeRowBytes = codeRowBytes; + this.numExPlanes = numExPlanes; + this.includeVectorNorm = includeVectorNorm; + } + + public PostingBlockBuilder addRow(String recordKey, + byte[] signPlane, + byte[] exPlanes, + float fAdd1, + float fRescale1, + float err1, + float fAddEx, + float fRescaleEx, + float residualNorm, + String fileGroupId, + String instantTime, + String partitionPath, + long rowPosition) { + return addRow( + recordKey, + signPlane, + exPlanes, + fAdd1, + fRescale1, + err1, + fAddEx, + fRescaleEx, + residualNorm, + null, + fileGroupId, + instantTime, + partitionPath, + rowPosition); + } + + public PostingBlockBuilder addRow(String recordKey, + byte[] signPlane, + byte[] exPlanes, + float fAdd1, + float fRescale1, + float err1, + float fAddEx, + float fRescaleEx, + float residualNorm, + Float vectorNorm, + String fileGroupId, + String instantTime, + String partitionPath, + long rowPosition) { + checkArgument(recordKey != null, "recordKey must not be null"); + checkArgument(signPlane != null && signPlane.length == codeRowBytes, + "signPlane must be exactly codeRowBytes"); + checkArgument(exPlanes != null && exPlanes.length == numExPlanes * codeRowBytes, + "exPlanes must be numExPlanes * codeRowBytes"); + checkArgument(rowPosition >= 0 && rowPosition <= 0xFFFFFFFFL, + "rowPosition must fit in unsigned int: " + rowPosition); + checkArgument(includeVectorNorm == (vectorNorm != null), + "vectorNorm presence must match block scalar layout"); + + int fileGroupIdx = fileGroupDictionaryIndex(fileGroupId, partitionPath); + int instantIdx = dictionaryIndex(instantTimeDict, instantTime); + rows.add(new Row( + recordKey, + signPlane.clone(), + exPlanes.clone(), + fAdd1, + fRescale1, + err1, + fAddEx, + fRescaleEx, + residualNorm, + vectorNorm == null ? 0.0f : vectorNorm, + fileGroupIdx, + instantIdx, + rowPosition)); + return this; + } + + public HoodieVectorIndexPostingBlock build() { + int numVectors = rows.size(); + ByteBuffer signPlane = allocateLittleEndian(numVectors * codeRowBytes); + ByteBuffer exPlanes = allocateLittleEndian(numVectors * numExPlanes * codeRowBytes); + ByteBuffer scalarFactors = allocateLittleEndian(numVectors * scalarFactorCount() * Float.BYTES); + ByteBuffer rowLocators = allocateLittleEndian(numVectors * ROW_LOCATOR_BYTES); + ByteBuffer recordKeyOffsets = allocateLittleEndian((numVectors + 1) * Integer.BYTES); + ByteBuffer recordKeyBytes = allocateLittleEndian(totalRecordKeyBytes()); + + for (Row row : rows) { + signPlane.put(row.signPlane); + exPlanes.put(row.exPlanes); + } + + writeScalarArray(scalarFactors, ScalarFactor.F_ADD_1); + writeScalarArray(scalarFactors, ScalarFactor.F_RESCALE_1); + writeScalarArray(scalarFactors, ScalarFactor.ERR_1); + writeScalarArray(scalarFactors, ScalarFactor.F_ADD_EX); + writeScalarArray(scalarFactors, ScalarFactor.F_RESCALE_EX); + writeScalarArray(scalarFactors, ScalarFactor.RESIDUAL_NORM); + if (includeVectorNorm) { + writeScalarArray(scalarFactors, ScalarFactor.VECTOR_NORM); + } + + int currentKeyOffset = 0; + recordKeyOffsets.putInt(currentKeyOffset); + for (Row row : rows) { + rowLocators.putShort((short) row.fileGroupIdx); + rowLocators.putShort((short) row.instantIdx); + rowLocators.putInt((int) row.rowPosition); + + byte[] keyBytes = row.recordKey.getBytes(StandardCharsets.UTF_8); + recordKeyBytes.put(keyBytes); + currentKeyOffset += keyBytes.length; + recordKeyOffsets.putInt(currentKeyOffset); + } + + return new HoodieVectorIndexPostingBlock( + BLOCK_FORMAT_VERSION, + numVectors, + codeRowBytes, + flip(signPlane), + flip(exPlanes), + flip(scalarFactors), + flip(rowLocators), + new ArrayList<>(fileGroupDict.keySet()), + new ArrayList<>(instantTimeDict.keySet()), + new ArrayList<>(partitionDict), + flip(recordKeyOffsets), + flip(recordKeyBytes)); + } + + public int rowCount() { + return rows.size(); + } + + public void reset() { + rows.clear(); + fileGroupDict.clear(); + instantTimeDict.clear(); + partitionDict.clear(); + } + + public static int deriveVectorsPerBlock(int targetBlockBytes, int dimPadded, int bitsTotal, int avgKeyLen) { + return deriveVectorsPerBlock(targetBlockBytes, dimPadded, bitsTotal, avgKeyLen, false); + } + + public static int deriveVectorsPerBlock(int targetBlockBytes, int dimPadded, int bitsTotal, int avgKeyLen, boolean includeVectorNorm) { + checkArgument(targetBlockBytes > 0, "targetBlockBytes must be positive"); + checkArgument(dimPadded > 0, "dimPadded must be positive"); + checkArgument(bitsTotal > 0, "bitsTotal must be positive"); + int codeRowBytes = ((dimPadded + 63) / 64) * Long.BYTES; + int scalarFactorCount = includeVectorNorm ? SCALAR_FACTOR_COUNT_WITH_VECTOR_NORM : SCALAR_FACTOR_COUNT; + int perVectorBytes = bitsTotal * codeRowBytes + scalarFactorCount * Float.BYTES + + ROW_LOCATOR_BYTES + Math.max(0, avgKeyLen) + Integer.BYTES; + int raw = targetBlockBytes / perVectorBytes; + int prevPowerOfTwo = Integer.highestOneBit(Math.max(1, raw)); + return Math.max(256, Math.min(4096, prevPowerOfTwo)); + } + + private void writeScalarArray(ByteBuffer scalarFactors, ScalarFactor factor) { + for (Row row : rows) { + scalarFactors.putFloat(row.scalar(factor)); + } + } + + private int scalarFactorCount() { + return includeVectorNorm ? SCALAR_FACTOR_COUNT_WITH_VECTOR_NORM : SCALAR_FACTOR_COUNT; + } + + private int totalRecordKeyBytes() { + int total = 0; + for (Row row : rows) { + total += row.recordKey.getBytes(StandardCharsets.UTF_8).length; + } + return total; + } + + private static int dictionaryIndex(Map dictionary, String value) { + checkArgument(value != null, "dictionary value must not be null"); + Integer existing = dictionary.get(value); + if (existing != null) { + return existing; + } + int index = dictionary.size(); + checkArgument(index <= 0xFFFF, "block-local dictionary cannot exceed unsigned short cardinality"); + dictionary.put(value, index); + return index; + } + + private int fileGroupDictionaryIndex(String fileGroupId, String partitionPath) { + checkArgument(fileGroupId != null, "fileGroupId must not be null"); + checkArgument(partitionPath != null, "partitionPath must not be null"); + Integer existing = fileGroupDict.get(fileGroupId); + if (existing != null) { + checkArgument(partitionDict.get(existing).equals(partitionPath), + "fileGroupId cannot map to multiple partition paths in one block: " + fileGroupId); + return existing; + } + int index = fileGroupDict.size(); + checkArgument(index <= 0xFFFF, "block-local dictionary cannot exceed unsigned short cardinality"); + fileGroupDict.put(fileGroupId, index); + partitionDict.add(partitionPath); + return index; + } + + private static ByteBuffer allocateLittleEndian(int size) { + return ByteBuffer.allocate(size).order(ByteOrder.LITTLE_ENDIAN); + } + + private static ByteBuffer flip(ByteBuffer buffer) { + buffer.flip(); + return buffer; + } + + private enum ScalarFactor { + F_ADD_1, + F_RESCALE_1, + ERR_1, + F_ADD_EX, + F_RESCALE_EX, + RESIDUAL_NORM, + VECTOR_NORM + } + + private static final class Row { + private final String recordKey; + private final byte[] signPlane; + private final byte[] exPlanes; + private final float fAdd1; + private final float fRescale1; + private final float err1; + private final float fAddEx; + private final float fRescaleEx; + private final float residualNorm; + private final float vectorNorm; + private final int fileGroupIdx; + private final int instantIdx; + private final long rowPosition; + + private Row(String recordKey, + byte[] signPlane, + byte[] exPlanes, + float fAdd1, + float fRescale1, + float err1, + float fAddEx, + float fRescaleEx, + float residualNorm, + float vectorNorm, + int fileGroupIdx, + int instantIdx, + long rowPosition) { + this.recordKey = recordKey; + this.signPlane = signPlane; + this.exPlanes = exPlanes; + this.fAdd1 = fAdd1; + this.fRescale1 = fRescale1; + this.err1 = err1; + this.fAddEx = fAddEx; + this.fRescaleEx = fRescaleEx; + this.residualNorm = residualNorm; + this.vectorNorm = vectorNorm; + this.fileGroupIdx = fileGroupIdx; + this.instantIdx = instantIdx; + this.rowPosition = rowPosition; + } + + private float scalar(ScalarFactor factor) { + switch (factor) { + case F_ADD_1: + return fAdd1; + case F_RESCALE_1: + return fRescale1; + case ERR_1: + return err1; + case F_ADD_EX: + return fAddEx; + case F_RESCALE_EX: + return fRescaleEx; + case RESIDUAL_NORM: + return residualNorm; + case VECTOR_NORM: + return vectorNorm; + default: + throw new IllegalArgumentException("Unknown scalar factor: " + factor); + } + } + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/PostingBlockView.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/PostingBlockView.java new file mode 100644 index 0000000000000..069a96acc458d --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/PostingBlockView.java @@ -0,0 +1,277 @@ +/* + * 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.hudi.common.index.vector; + +import org.apache.hudi.avro.model.HoodieVectorIndexPostingBlock; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.util.List; + +import static org.apache.hudi.common.util.ValidationUtils.checkArgument; + +/** + * Zero-copy section view over a vector posting block payload. + */ +public final class PostingBlockView { + + private final HoodieVectorIndexPostingBlock block; + private final int numVectors; + private final int codeRowBytes; + private final int numExPlanes; + private final ByteBuffer signPlane; + private final ByteBuffer exPlanes; + private final ByteBuffer scalarFactors; + private final int scalarFactorCount; + private final ByteBuffer rowLocators; + private final ByteBuffer recordKeyOffsets; + private final ByteBuffer recordKeyBytes; + + public PostingBlockView(HoodieVectorIndexPostingBlock block) { + checkArgument(block != null, "posting block must not be null"); + this.block = block; + this.numVectors = block.getNumVectors(); + this.codeRowBytes = block.getCodeRowBytes(); + this.signPlane = littleEndian(block.getSignPlane()); + this.exPlanes = littleEndian(block.getExPlanes()); + this.scalarFactors = littleEndian(block.getScalarFactors()); + this.rowLocators = littleEndian(block.getRowLocators()); + this.recordKeyOffsets = littleEndian(block.getRecordKeyOffsets()); + this.recordKeyBytes = littleEndian(block.getRecordKeyBytes()); + + checkArgument(block.getBlockFormatVersion() == PostingBlockBuilder.BLOCK_FORMAT_VERSION, + "Unsupported posting block format version: " + block.getBlockFormatVersion()); + checkArgument(numVectors >= 0, "numVectors must be non-negative"); + checkArgument(codeRowBytes > 0 && codeRowBytes % Long.BYTES == 0, + "codeRowBytes must be positive and long-aligned"); + checkArgument(signPlane.remaining() == numVectors * codeRowBytes, + "signPlane length does not match numVectors * codeRowBytes"); + checkArgument(exPlanes.remaining() % Math.max(1, numVectors * codeRowBytes) == 0, + "exPlanes length is not row aligned"); + this.numExPlanes = numVectors == 0 ? 0 : exPlanes.remaining() / (numVectors * codeRowBytes); + int scalarArrayBytes = numVectors * Float.BYTES; + checkArgument(numVectors == 0 || scalarFactors.remaining() % scalarArrayBytes == 0, + "scalarFactors length must be a whole number of float arrays"); + this.scalarFactorCount = numVectors == 0 + ? PostingBlockBuilder.SCALAR_FACTOR_COUNT + : scalarFactors.remaining() / scalarArrayBytes; + checkArgument(scalarFactorCount == PostingBlockBuilder.SCALAR_FACTOR_COUNT + || scalarFactorCount == PostingBlockBuilder.SCALAR_FACTOR_COUNT_WITH_VECTOR_NORM, + "scalarFactors length must contain six or seven float arrays"); + checkArgument(rowLocators.remaining() == numVectors * PostingBlockBuilder.ROW_LOCATOR_BYTES, + "rowLocators length must use 8-byte locator stride"); + checkArgument(recordKeyOffsets.remaining() == (numVectors + 1) * Integer.BYTES, + "recordKeyOffsets length must be (numVectors + 1) * 4"); + } + + public int numVectors() { + return numVectors; + } + + public int codeRowBytes() { + return codeRowBytes; + } + + public int numExPlanes() { + return numExPlanes; + } + + public boolean hasVectorNorm() { + return scalarFactorCount == PostingBlockBuilder.SCALAR_FACTOR_COUNT_WITH_VECTOR_NORM; + } + + public int signPlaneOffset(int vectorIndex) { + checkVectorIndex(vectorIndex); + return vectorIndex * codeRowBytes; + } + + public int exPlaneOffset(int vectorIndex, int exPlaneIndex) { + checkVectorIndex(vectorIndex); + checkArgument(exPlaneIndex >= 0 && exPlaneIndex < numExPlanes, + "exPlaneIndex out of bounds: " + exPlaneIndex); + return vectorIndex * numExPlanes * codeRowBytes + exPlaneIndex * codeRowBytes; + } + + public ByteBuffer signPlaneRow(int vectorIndex) { + return slice(signPlane, signPlaneOffset(vectorIndex), codeRowBytes); + } + + /** + * Returns a duplicate of the whole sign-plane section. Hot scan loops should + * use absolute {@code get(offset)} reads from this buffer instead of slicing + * per vector. + */ + public ByteBuffer signPlaneBuffer() { + return signPlane.duplicate().order(ByteOrder.LITTLE_ENDIAN); + } + + /** + * Returns a duplicate of the whole extended-plane section for survivor-only + * repacking. + */ + public ByteBuffer exPlanesBuffer() { + return exPlanes.duplicate().order(ByteOrder.LITTLE_ENDIAN); + } + + public ByteBuffer exPlaneRow(int vectorIndex, int exPlaneIndex) { + return slice(exPlanes, exPlaneOffset(vectorIndex, exPlaneIndex), codeRowBytes); + } + + public float scalarFactor(ScalarFactor factor, int vectorIndex) { + checkVectorIndex(vectorIndex); + checkScalarFactor(factor); + int offset = scalarFactorOffset(factor, vectorIndex); + return scalarFactors.getFloat(offset); + } + + public float vectorNormOrNaN(int vectorIndex) { + return hasVectorNorm() ? scalarFactor(ScalarFactor.VECTOR_NORM, vectorIndex) : Float.NaN; + } + + public int scalarFactorOffset(ScalarFactor factor, int vectorIndex) { + checkVectorIndex(vectorIndex); + checkScalarFactor(factor); + return factor.ordinal() * numVectors * Float.BYTES + vectorIndex * Float.BYTES; + } + + public RowLocator rowLocator(int vectorIndex) { + checkVectorIndex(vectorIndex); + int offset = vectorIndex * PostingBlockBuilder.ROW_LOCATOR_BYTES; + int fileGroupIdx = Short.toUnsignedInt(rowLocators.getShort(offset)); + int instantIdx = Short.toUnsignedInt(rowLocators.getShort(offset + Short.BYTES)); + long rowPosition = Integer.toUnsignedLong(rowLocators.getInt(offset + 2 * Short.BYTES)); + return new RowLocator( + fileGroupIdx, + instantIdx, + rowPosition, + block.getFileGroupDict().get(fileGroupIdx), + block.getInstantTimeDict().get(instantIdx), + partitionForFileGroup(fileGroupIdx)); + } + + public String recordKey(int vectorIndex) { + checkVectorIndex(vectorIndex); + int start = recordKeyOffsets.getInt(vectorIndex * Integer.BYTES); + int end = recordKeyOffsets.getInt((vectorIndex + 1) * Integer.BYTES); + checkArgument(start >= 0 && end >= start && end <= recordKeyBytes.remaining(), + "Invalid record key offset range"); + ByteBuffer keySlice = slice(recordKeyBytes, start, end - start); + byte[] bytes = new byte[keySlice.remaining()]; + keySlice.get(bytes); + return new String(bytes, StandardCharsets.UTF_8); + } + + public List fileGroupDict() { + return block.getFileGroupDict(); + } + + public List instantTimeDict() { + return block.getInstantTimeDict(); + } + + public List partitionDict() { + return block.getPartitionDict(); + } + + private String partitionForFileGroup(int fileGroupIdx) { + return block.getPartitionDict().isEmpty() ? "" : block.getPartitionDict().get(fileGroupIdx); + } + + private void checkVectorIndex(int vectorIndex) { + checkArgument(vectorIndex >= 0 && vectorIndex < numVectors, "vectorIndex out of bounds: " + vectorIndex); + } + + private void checkScalarFactor(ScalarFactor factor) { + checkArgument(factor.ordinal() < scalarFactorCount, + "Scalar factor is absent from this block layout: " + factor); + } + + private static ByteBuffer littleEndian(ByteBuffer buffer) { + ByteBuffer duplicate = buffer.duplicate(); + duplicate.order(ByteOrder.LITTLE_ENDIAN); + return duplicate; + } + + private static ByteBuffer slice(ByteBuffer buffer, int offset, int length) { + ByteBuffer duplicate = buffer.duplicate(); + duplicate.position(offset); + duplicate.limit(offset + length); + ByteBuffer slice = duplicate.slice(); + slice.order(ByteOrder.LITTLE_ENDIAN); + return slice; + } + + public enum ScalarFactor { + F_ADD_1, + F_RESCALE_1, + ERR_1, + F_ADD_EX, + F_RESCALE_EX, + RESIDUAL_NORM, + VECTOR_NORM + } + + public static final class RowLocator { + private final int fileGroupDictIndex; + private final int instantTimeDictIndex; + private final long rowPosition; + private final String fileGroupId; + private final String instantTime; + private final String partitionPath; + + private RowLocator(int fileGroupDictIndex, + int instantTimeDictIndex, + long rowPosition, + String fileGroupId, + String instantTime, + String partitionPath) { + this.fileGroupDictIndex = fileGroupDictIndex; + this.instantTimeDictIndex = instantTimeDictIndex; + this.rowPosition = rowPosition; + this.fileGroupId = fileGroupId; + this.instantTime = instantTime; + this.partitionPath = partitionPath; + } + + public int getFileGroupDictIndex() { + return fileGroupDictIndex; + } + + public int getInstantTimeDictIndex() { + return instantTimeDictIndex; + } + + public long getRowPosition() { + return rowPosition; + } + + public String getFileGroupId() { + return fileGroupId; + } + + public String getInstantTime() { + return instantTime; + } + + public String getPartitionPath() { + return partitionPath; + } + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/QuantizedVector.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/QuantizedVector.java new file mode 100644 index 0000000000000..962952990aa3e --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/QuantizedVector.java @@ -0,0 +1,169 @@ +/* + * 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.hudi.common.index.vector; + +import java.io.Serializable; + +/** Encoded vector and its persisted RaBitQ scoring factors. */ +public final class QuantizedVector implements Serializable { + + private static final long serialVersionUID = 1L; + + /** Primary packed sign/MSB code. */ + final byte[] code; + /** Optional packed lower-bit code planes for multibit RaBitQ. */ + final byte[] extendedCode; + /** Original or residual vector norm. */ + final float scalar; + /** Optional L2-estimator additive factor used by multibit MDT postings. */ + final Float additiveFactor; + /** Optional L2-estimator rescale factor used by multibit MDT postings. */ + final Float rescaleFactor; + /** Sign-only L2-estimator additive factor used by packed-block pass 1. */ + final Float additiveFactor1; + /** Sign-only L2-estimator rescale factor used by packed-block pass 1. */ + final Float rescaleFactor1; + /** Query-independent pass-1 distance error factor. */ + final Float error1; + /** Optional raw vector norm for metric-neutral residual scoring. */ + final Float vectorNorm; + /** Total RaBitQ bit width represented by this payload. */ + final int bits; + + public QuantizedVector(byte[] code, float scalar) { + this(code, null, scalar, null, null, null, null, null, null, 1); + } + + public QuantizedVector(byte[] code, + byte[] extendedCode, + float scalar, + Float additiveFactor, + Float rescaleFactor, + int bits) { + this(code, extendedCode, scalar, additiveFactor, rescaleFactor, null, null, null, null, bits); + } + + public QuantizedVector(byte[] code, + byte[] extendedCode, + float scalar, + Float additiveFactor, + Float rescaleFactor, + Float additiveFactor1, + Float rescaleFactor1, + Float error1, + int bits) { + this(code, extendedCode, scalar, additiveFactor, rescaleFactor, + additiveFactor1, rescaleFactor1, error1, null, bits); + } + + public QuantizedVector(byte[] code, + byte[] extendedCode, + float scalar, + Float additiveFactor, + Float rescaleFactor, + Float additiveFactor1, + Float rescaleFactor1, + Float error1, + Float vectorNorm, + int bits) { + if (code == null || code.length == 0) { + throw new IllegalArgumentException("Primary code must not be empty"); + } + if (bits < 1 || bits > 8) { + throw new IllegalArgumentException("RaBitQ bits must be in [1, 8]: " + bits); + } + if (!Float.isFinite(scalar) || scalar < 0f) { + throw new IllegalArgumentException("Scalar must be finite and non-negative: " + scalar); + } + if ((bits == 1 && extendedCode != null && extendedCode.length != 0) + || (bits > 1 && (extendedCode == null || extendedCode.length == 0))) { + throw new IllegalArgumentException("Extended code presence must match the RaBitQ bit width"); + } + validateOptionalFactor(additiveFactor, "additiveFactor"); + validateNonNegativeFactor(rescaleFactor, "rescaleFactor"); + validateOptionalFactor(additiveFactor1, "additiveFactor1"); + validateNonNegativeFactor(rescaleFactor1, "rescaleFactor1"); + validateNonNegativeFactor(error1, "error1"); + validateNonNegativeFactor(vectorNorm, "vectorNorm"); + this.code = code.clone(); + this.extendedCode = extendedCode == null ? null : extendedCode.clone(); + this.scalar = scalar; + this.additiveFactor = additiveFactor; + this.rescaleFactor = rescaleFactor; + this.additiveFactor1 = additiveFactor1; + this.rescaleFactor1 = rescaleFactor1; + this.error1 = error1; + this.vectorNorm = vectorNorm; + this.bits = bits; + } + + public byte[] getCode() { + return code.clone(); + } + + public byte[] getExtendedCode() { + return extendedCode == null ? null : extendedCode.clone(); + } + + public float getScalar() { + return scalar; + } + + public Float getAdditiveFactor() { + return additiveFactor; + } + + public Float getRescaleFactor() { + return rescaleFactor; + } + + public Float getAdditiveFactor1() { + return additiveFactor1; + } + + public Float getRescaleFactor1() { + return rescaleFactor1; + } + + public Float getError1() { + return error1; + } + + public Float getVectorNorm() { + return vectorNorm; + } + + public int getBits() { + return bits; + } + + private static void validateOptionalFactor(Float value, String name) { + if (value != null && !Float.isFinite(value)) { + throw new IllegalArgumentException(name + " must be finite: " + value); + } + } + + private static void validateNonNegativeFactor(Float value, String name) { + validateOptionalFactor(value, name); + if (value != null && value < 0f) { + throw new IllegalArgumentException(name + " must be non-negative: " + value); + } + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQByteLutScorer.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQByteLutScorer.java new file mode 100644 index 0000000000000..6454351fbcbc0 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQByteLutScorer.java @@ -0,0 +1,169 @@ +/* + * 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.hudi.common.index.vector; + +import java.nio.ByteBuffer; + +/** + * Byte-lookup-table posting scorer (RFC-109 §3A, fixes 1 and 2). + * + *

Exact float-query semantics. This scorer scores the same float rotated (residual) + * query used by {@link RaBitQEncoder#dotPackedBinary} and {@link RaBitQEncoder#multibitDotTerm}. + * It introduces no query quantization (unlike {@link RaBitQPlaneKernel}, which quantizes + * the query to {@code Bq} planes and therefore changes recall math). It only re-associates the + * per-dimension sum {@code dot(q, code)} into per-byte partial sums precomputed once per probed + * cluster, which lets the scan replace: + *

+ * with table lookups and zero per-survivor allocation. Results match the scalar path up to + * floating-point re-association (byte grouping), never a semantic (quantization) change. + * + *

LUT layout. {@code lut[bytePos][pattern]} holds the sum of {@code query[bytePos*8 + b]} + * over the set bits {@code b} of {@code pattern}, with padding dimensions ({@code >= dimension}) + * contributing zero so the byte-grouped sum equals the {@code [0, dimension)} scalar sum. A packed + * plane dot is then {@code sum(lut[bytePos][planeByte])} over byte positions — {@code codeRowBytes} + * lookups instead of {@code dimension} branchy bit tests (an ~8x op reduction at any dimension). + */ +public final class RaBitQByteLutScorer { + + private final double[][] lut; // [codeRowBytes][256] + private final int codeRowBytes; + private final float querySum; + + private RaBitQByteLutScorer(double[][] lut, int codeRowBytes, float querySum) { + this.lut = lut; + this.codeRowBytes = codeRowBytes; + this.querySum = querySum; + } + + /** + * Build the per-cluster LUT from the (residual) rotated query. Called at most once per distinct + * probed cluster; build cost is {@code codeRowBytes * 256} adds, amortized over the cluster's + * posting scan. + * + * @param rotatedQuery the rotated residual query {@code wRot} (length {@code >= dimension}) + * @param querySum {@code sum(wRot)}; folded into the pass-1/pass-2 centering terms + * @param dimension raw dimension scored (padding dims contribute zero) + * @param codeRowBytes long-aligned per-plane row width from the block layout + */ + public static RaBitQByteLutScorer forQuery(float[] rotatedQuery, float querySum, + int dimension, int codeRowBytes) { + if (rotatedQuery == null || dimension <= 0 || rotatedQuery.length < dimension) { + throw new IllegalArgumentException("rotatedQuery must contain every scored dimension"); + } + if (!Float.isFinite(querySum) || codeRowBytes < (dimension + Byte.SIZE - 1) / Byte.SIZE) { + throw new IllegalArgumentException("querySum and codeRowBytes must match the scored query"); + } + for (int i = 0; i < dimension; i++) { + if (!Float.isFinite(rotatedQuery[i])) { + throw new IllegalArgumentException("rotatedQuery contains a non-finite value at dimension " + i); + } + } + double[][] lut = new double[codeRowBytes][256]; + for (int bytePos = 0; bytePos < codeRowBytes; bytePos++) { + int baseDim = bytePos << 3; + double[] bitContribution = new double[8]; + for (int b = 0; b < 8; b++) { + int dim = baseDim + b; + bitContribution[b] = dim < dimension ? rotatedQuery[dim] : 0.0; + } + double[] table = lut[bytePos]; + for (int pattern = 0; pattern < 256; pattern++) { + double sum = 0.0; + // Ascending bit order mirrors dotSignRow's ascending-dimension accumulation. + for (int b = 0; b < 8; b++) { + if ((pattern & (1 << b)) != 0) { + sum += bitContribution[b]; + } + } + table[pattern] = sum; + } + } + return new RaBitQByteLutScorer(lut, codeRowBytes, querySum); + } + + /** + * Raw packed-plane inner product {@code dot(query, plane)} read directly from a plane buffer at an + * absolute byte offset. Equivalent to {@link RaBitQEncoder#dotPackedBinary} for the sign plane. + */ + public double planeDot(ByteBuffer planeBuffer, int offset) { + if (planeBuffer == null || offset < 0 || offset > planeBuffer.limit() - codeRowBytes) { + throw new IllegalArgumentException("Plane row exceeds the supplied buffer"); + } + double sum = 0.0; + for (int bytePos = 0; bytePos < codeRowBytes; bytePos++) { + sum += lut[bytePos][planeBuffer.get(offset + bytePos) & 0xFF]; + } + return sum; + } + + /** + * Pass-1 sign-only score {@code dot(query, sign) - 0.5*sumQuery} (== {@code dotSignRow}). Callers + * that also run pass-2 should keep the {@link #planeDot} sign value and reuse it via + * {@link #pass1FromDot(double)} and {@link #pass2(double, PostingBlockView, ByteBuffer, int, int, int)} + * rather than recomputing the sign dot. + */ + public float pass1(ByteBuffer signBuffer, int signOffset) { + return pass1FromDot(planeDot(signBuffer, signOffset)); + } + + /** Pass-1 score from an already-computed sign-plane dot (see {@link #planeDot}). */ + public float pass1FromDot(double signDot) { + return (float) (signDot + querySum * -0.5f); + } + + /** + * Pass-2 full multibit dot term, scored directly from the sign plane and extended bit-planes + * with no repacking and no per-survivor allocation. Bit-plane decomposition of the centered + * code makes this algebraically identical to {@link RaBitQEncoder#multibitDotTerm}: + * 

+   *   extendedDot = sum_p 2^(exBits-1-p) * dot(query, exPlane_p)
+   *   dotTerm     = 2^exBits * signDot + extendedDot + sumQuery * -((2^bits - 1)/2)
+   *
+ * + * @param signDot the raw sign-plane dot from {@link #planeDot} (reuse the pass-1 value) + * @param view the posting block view (for extended-plane offsets) + * @param exBuffer the extended-planes buffer ({@link PostingBlockView#exPlanesBuffer()}) + * @param vectorIndex the vector ordinal within the block + * @param exBits number of extended planes ({@code bits - 1}) + * @param bits total RaBitQ bits + */ + public float pass2(double signDot, PostingBlockView view, ByteBuffer exBuffer, + int vectorIndex, int exBits, int bits) { + if (view == null || exBuffer == null || bits < 1 || bits > 8 || exBits != bits - 1) { + throw new IllegalArgumentException("Posting view and a consistent 1-8 bit width are required"); + } + if (view.codeRowBytes() != codeRowBytes || view.numExPlanes() != exBits) { + throw new IllegalArgumentException("Scorer and posting block layouts do not match"); + } + view.signPlaneOffset(vectorIndex); // validates the vector index before any early return + if (exBits <= 0) { + return (float) (signDot + querySum * -0.5f); + } + double extendedDot = 0.0; + for (int p = 0; p < exBits; p++) { + extendedDot += (double) (1L << (exBits - 1 - p)) * planeDot(exBuffer, view.exPlaneOffset(vectorIndex, p)); + } + double cBias = -((double) ((1 << bits) - 1)) / 2.0; + return (float) (((double) (1 << exBits)) * signDot + extendedDot + querySum * cBias); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQDistanceScorer.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQDistanceScorer.java new file mode 100644 index 0000000000000..22c6d4b79ef03 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQDistanceScorer.java @@ -0,0 +1,126 @@ +/* + * 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.hudi.common.index.vector; + +/** Distance reconstruction for sign-encoded RaBitQ vectors. */ +public final class RaBitQDistanceScorer { + + private RaBitQDistanceScorer() { + } + + /** Estimates the legacy norm-weighted angular proxy from two packed sign codes. */ + public static float estimateDistance(byte[] queryCode, byte[] encodedCode, + float scalar, int dimension) { + if (!Float.isFinite(scalar) || scalar < 0f) { + throw new IllegalArgumentException("scalar must be finite and non-negative: " + scalar); + } + float cosine = symmetricCosine(queryCode, encodedCode, dimension); + return scalar * (1.0f - cosine); + } + + /** Cosine estimate from binary-vs-binary Hamming distance. */ + public static float symmetricCosine(byte[] queryCode, byte[] encodedCode, int dimension) { + validatePackedCode(queryCode, dimension, "queryCode"); + validatePackedCode(encodedCode, dimension, "encodedCode"); + int fullBytes = dimension / Byte.SIZE; + int hamming = 0; + for (int i = 0; i < fullBytes; i++) { + hamming += Integer.bitCount((queryCode[i] ^ encodedCode[i]) & 0xFF); + } + int remainingBits = dimension % Byte.SIZE; + if (remainingBits != 0) { + int mask = (1 << remainingBits) - 1; + hamming += Integer.bitCount((queryCode[fullBytes] ^ encodedCode[fullBytes]) & mask); + } + return 1.0f - 2.0f * hamming / dimension; + } + + /** + * Cosine estimate from the rotated float query and a packed sign code. + * This retains full query precision and therefore has lower variance than Hamming scoring. + */ + public static float asymmetricCosine(float[] rotatedQuery, byte[] encodedCode, int dimension) { + if (rotatedQuery == null || rotatedQuery.length != dimension) { + throw new IllegalArgumentException("rotatedQuery length must equal dimension " + dimension); + } + validatePackedCode(encodedCode, dimension, "encodedCode"); + double sum = 0.0; + for (int i = 0; i < dimension; i++) { + if (!Float.isFinite(rotatedQuery[i])) { + throw new IllegalArgumentException("rotatedQuery contains a non-finite value at dimension " + i); + } + boolean positive = (encodedCode[i >> 3] & (1 << (i & 7))) != 0; + sum += positive ? rotatedQuery[i] : -rotatedQuery[i]; + } + return (float) (sum / Math.sqrt(dimension)); + } + + /** Reconstructs the configured distance from a cosine estimate and vector norms. */ + public static float reconstructDistance(VectorDistanceMetric metric, + float cosineEstimate, + float queryNorm, + float vectorNorm) { + if (metric == null || !Float.isFinite(cosineEstimate) + || !Float.isFinite(queryNorm) || queryNorm < 0f + || !Float.isFinite(vectorNorm) || vectorNorm < 0f) { + throw new IllegalArgumentException("Metric, cosine estimate, and norms must be valid"); + } + float cosine = Math.max(-1.0f, Math.min(1.0f, cosineEstimate)); + switch (metric) { + case L2: + double squaredDistance = (double) queryNorm * queryNorm + + (double) vectorNorm * vectorNorm + - 2.0 * queryNorm * vectorNorm * cosine; + return (float) Math.sqrt(Math.max(0.0, squaredDistance)); + case DOT_PRODUCT: + return -(queryNorm * vectorNorm * cosine); + case COSINE: + default: + return queryNorm == 0f || vectorNorm == 0f ? 1.0f : 1.0f - cosine; + } + } + + /** Computes Hamming distance between equally sized packed binary codes. */ + public static int hammingDistance(byte[] left, byte[] right) { + if (left == null || right == null) { + throw new IllegalArgumentException("Packed codes must not be null"); + } + if (left.length != right.length) { + throw new IllegalArgumentException( + "Packed code length mismatch: " + left.length + " != " + right.length); + } + int count = 0; + for (int i = 0; i < left.length; i++) { + count += Integer.bitCount((left[i] ^ right[i]) & 0xFF); + } + return count; + } + + private static void validatePackedCode(byte[] code, int dimension, String name) { + if (dimension <= 0) { + throw new IllegalArgumentException("dimension must be positive: " + dimension); + } + int expectedBytes = (dimension + Byte.SIZE - 1) / Byte.SIZE; + if (code == null || code.length != expectedBytes) { + throw new IllegalArgumentException( + name + " length must be " + expectedBytes + " bytes for dimension " + dimension); + } + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQEncoder.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQEncoder.java new file mode 100644 index 0000000000000..cc5ef0db2f6ca --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQEncoder.java @@ -0,0 +1,588 @@ +/* + * 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.hudi.common.index.vector; + +import java.io.Serializable; +import java.util.Map; +import java.util.Objects; +import java.util.Random; +import java.util.concurrent.ConcurrentHashMap; + +/** + * RaBitQ (Randomized Binary Quantization) encoder. + * + *

Encodes float vectors into 1-bit-per-dimension binary codes using a random + * orthogonal rotation matrix R. The rotation is deterministic given (seed, dimension), + * so R is stored as a seed — it never needs retraining. + * + *

Encoding (write path): + *

    + *
  1. Normalize: v̂ = v / ||v||
    2. + *
  2. Rotate: v_rot = R @ v̂
    3. + *
  3. Binarize: code = pack(sign(v_rot)) — D bits → ceil(D/8) bytes
    4. + *
  4. Scalar: s = ||v|| (stored alongside code; 1.0 if assume_normalized)
    5. + *
+ * + *

Query scan (read path): + *

    + *
  1. Normalize + rotate the query vector (same R)
    2. + *
  2. Binarize to get q_bin
    3. + *
  3. For each database code b: Hamming(q_bin, b) → estimated cosine
    4. + *
  4. Re-rank top-R candidates with exact distance
    5. + *
+ * + *

Memory: the rotation matrix is D×D floats = 4·D² bytes (~2.3 MB at D=768). + * It is built lazily on first use and reused across encode calls. + * + *

Thread-safe after the first call to {@link #encode} or {@link #encodeQuery} + * (the lazy init is synchronized). + */ +public final class RaBitQEncoder implements Serializable { + private static final long serialVersionUID = 1L; + private static final Map ROTATION_MATRIX_CACHE = new ConcurrentHashMap<>(); + /** Neutral-factor thresholds for the current posting-block format (RFC-109 §3). */ + private static final RaBitQFactorConfig FACTOR_CONFIG = RaBitQFactorConfig.defaults(); + private final int dimension; + private final int bits; + private final long seed; + private final boolean assumeNormalized; + + /** Rotation matrix, row-major. Populated lazily. */ + private transient volatile float[][] rotMat; + public RaBitQEncoder(int dimension, long seed, boolean assumeNormalized) { + this(dimension, 1, seed, assumeNormalized); + } + + public RaBitQEncoder(int dimension, int bits, long seed, boolean assumeNormalized) { + if (dimension <= 0) { + throw new IllegalArgumentException("Dimension must be positive, got: " + dimension); + } + if (bits <= 0 || bits > 8) { + throw new IllegalArgumentException("RaBitQ bits must be in [1, 8], got: " + bits); + } + this.dimension = dimension; + this.bits = bits; + this.seed = seed; + this.assumeNormalized = assumeNormalized; + } + + /** Convenience constructor using default seed. */ + public RaBitQEncoder(int dimension) { + this(dimension, 1, 42L, false); + } + + // ---- encoding ---------------------------------------------------------- + + public QuantizedVector encode(float[] vector) { + validateVector(vector, "vector"); + float norm = norm(vector); + float scalar = assumeNormalized ? 1.0f : norm; + float[] normalized = (norm == 0f || assumeNormalized) + ? vector + : normalize(vector, norm); + float[] rotated = rotate(normalized); + return new QuantizedVector(binarize(rotated), scalar); + } + + /** + * Metric-neutral residual encoding used by MDT postings. + * + *

The optional {@code center} is the IVF centroid in the original vector space. + * When null, the vector is encoded relative to the origin. + */ + public QuantizedVector encodeResidual(float[] vector, float[] center) { + validateVector(vector, "vector"); + if (center != null) { + validateVector(center, "center"); + } + float[] rotatedVector = rotate(vector); + float[] rotatedCenter = center == null ? new float[dimension] : rotate(center); + float[] residual = subtract(rotatedVector, rotatedCenter); + byte[] binaryCode = binarize(residual); + + float residualNorm = norm(residual); + if (residualNorm == 0.0f) { + return new QuantizedVector( + binaryCode, + new byte[extendedCodeBytes()], + 0.0f, + 0.0f, + 0.0f, + 0.0f, + 0.0f, + 0.0f, + norm(vector), + bits); + } + + if (bits <= 1) { + double ipResidual1 = 0.0d; + for (int i = 0; i < dimension; i++) { + ipResidual1 += residual[i] * (residual[i] > 0f ? 0.5d : -0.5d); + } + RaBitQNeutralFactors.Factors factors = RaBitQNeutralFactors.compute( + residual, + rotatedCenter, + rotatedVector, + ipResidual1, + ipResidual1, + dimension, + FACTOR_CONFIG); + return new QuantizedVector( + binaryCode, + new byte[0], + factors.residualNorm, + factors.centerRip, + factors.fRescale1, + factors.centerRip, + factors.fRescale1, + factors.err1, + factors.vectorNorm, + bits); + } + + int exBits = bits - 1; + + float[] absNormalizedResidual = new float[dimension]; + for (int i = 0; i < dimension; i++) { + absNormalizedResidual[i] = Math.abs(residual[i] / residualNorm); + } + + QuantizedLevels quantizedLevels = quantizeEx(absNormalizedResidual, exBits); + int[] exCode = quantizedLevels.levels; + int mask = (1 << exBits) - 1; + for (int i = 0; i < dimension; i++) { + if (residual[i] < 0f) { + exCode[i] = mask - exCode[i]; + } + } + + float cBias = (float) -((1 << exBits) - 0.5d); + double ipResidual1 = 0.0d; + double ipResidual = 0.0d; + for (int i = 0; i < dimension; i++) { + boolean positive = residual[i] > 0f; + double signOnlyCode = positive ? 0.5d : -0.5d; + float centeredCode = exCode[i] + (positive ? (1 << exBits) : 0) + cBias; + ipResidual1 += residual[i] * signOnlyCode; + ipResidual += residual[i] * centeredCode; + } + + RaBitQNeutralFactors.Factors factors = RaBitQNeutralFactors.compute( + residual, + rotatedCenter, + rotatedVector, + ipResidual1, + ipResidual, + dimension, + FACTOR_CONFIG); + return new QuantizedVector( + binaryCode, + packUnsignedLevels(exCode, exBits), + factors.residualNorm, + factors.centerRip, + factors.fRescaleEx, + factors.centerRip, + factors.fRescale1, + factors.err1, + factors.vectorNorm, + bits); + } + + /** + * Compatibility shim for old call sites. New postings must use the metric-neutral + * residual factor convention produced by {@link #encodeResidual(float[], float[])}. + */ + public QuantizedVector encodeForL2(float[] vector, float[] center) { + return encodeResidual(vector, center); + } + + public RaBitQQueryState encodeQuery(float[] queryVector) { + validateVector(queryVector, "query"); + float norm = norm(queryVector); + float[] normalized = norm == 0f ? queryVector : normalize(queryVector, norm); + float[] rotated = rotate(normalized); + return new RaBitQQueryState(binarize(rotated), rotated, norm); + } + + public RaBitQQueryState encodeQueryForL2(float[] queryVector) { + validateVector(queryVector, "query"); + float[] rotated = rotate(queryVector); + return new RaBitQQueryState(binarize(rotated), rotated, norm(queryVector)); + } + + public float estimateDistance(RaBitQQueryState queryState, QuantizedVector encoded) { + return RaBitQDistanceScorer.estimateDistance( + queryState.binaryCodeUnsafe(), encoded.code, encoded.scalar, dimension); + } + + /** Metric-aware approximate distance using symmetric or asymmetric cosine estimation. */ + public float estimateDistance(RaBitQQueryState queryState, + QuantizedVector encoded, + VectorDistanceMetric metric, + boolean asymmetric) { + float cosine = asymmetric + ? RaBitQDistanceScorer.asymmetricCosine(queryState.rotatedQueryUnsafe(), encoded.code, dimension) + : RaBitQDistanceScorer.symmetricCosine(queryState.binaryCodeUnsafe(), encoded.code, dimension); + return RaBitQDistanceScorer.reconstructDistance( + metric, cosine, queryState.getQueryNorm(), encoded.scalar); + } + + public int codeBytes() { + return (dimension + 7) / 8; + } + + public int extendedCodeBytes() { + return bits <= 1 ? 0 : (dimension * (bits - 1) + 7) / 8; + } + + public int totalCodeBytes() { + return codeBytes() + extendedCodeBytes(); + } + + public int getBits() { + return bits; + } + + // ---- public utilities -------------------------------------------------- + + /** L2 norm of a float vector. */ + public static float norm(float[] v) { + double sum = 0.0; + for (float x : v) { + sum += (double) x * x; + } + return (float) Math.sqrt(sum); + } + + public static float l2Squared(float[] left, float[] right) { + double sum = 0.0; + for (int i = 0; i < left.length; i++) { + double delta = (double) left[i] - right[i]; + sum += delta * delta; + } + return (float) sum; + } + + public float[] rotateVector(float[] vector) { + validateVector(vector, "vector"); + return rotate(vector); + } + + public static float dotPackedBinary(byte[] binaryCode, float[] query, int dimension) { + double sum = 0.0; + for (int i = 0; i < dimension; i++) { + if ((binaryCode[i >> 3] & (1 << (i & 7))) != 0) { + sum += query[i]; + } + } + return (float) sum; + } + + public static float dotPackedUnsigned(byte[] packedLevels, int bitsPerValue, float[] query, int dimension) { + if (packedLevels == null || bitsPerValue <= 0) { + return 0.0f; + } + double sum = 0.0; + int bitOffset = 0; + for (int i = 0; i < dimension; i++) { + int value = 0; + for (int bit = 0; bit < bitsPerValue; bit++) { + int absoluteBit = bitOffset + bit; + int byteIndex = absoluteBit >> 3; + int bitIndex = absoluteBit & 7; + if ((packedLevels[byteIndex] & (1 << bitIndex)) != 0) { + value |= (1 << bit); + } + } + sum += (double) value * query[i]; + bitOffset += bitsPerValue; + } + return (float) sum; + } + + public static float multibitDotTerm(float[] rotatedQuery, + float sumQuery, + byte[] binaryCode, + byte[] extendedCode, + int dimension, + int bits) { + int exBits = bits - 1; + float binaryDot = dotPackedBinary(binaryCode, rotatedQuery, dimension); + if (exBits <= 0) { + return binaryDot + (sumQuery * -0.5f); + } + float extendedDot = dotPackedUnsigned(extendedCode, exBits, rotatedQuery, dimension); + float cBias = (float) -((1 << bits) - 1) / 2.0f; + return ((float) (1 << exBits) * binaryDot) + extendedDot + (sumQuery * cBias); + } + + public static float multibitDotTerm(float[] rotatedQuery, + byte[] binaryCode, + byte[] extendedCode, + int dimension, + int bits) { + return multibitDotTerm(rotatedQuery, sum(rotatedQuery), binaryCode, extendedCode, dimension, bits); + } + + // ---- private ----------------------------------------------------------- + + private void validateVector(float[] vector, String name) { + if (vector == null) { + throw new IllegalArgumentException(name + " must not be null"); + } + if (vector.length != dimension) { + throw new IllegalArgumentException( + "Expected " + name + " dimension " + dimension + ", got " + vector.length); + } + for (int i = 0; i < vector.length; i++) { + if (!Float.isFinite(vector[i])) { + throw new IllegalArgumentException( + name + " contains a non-finite value at dimension " + i + ": " + vector[i]); + } + } + } + + private float[] normalize(float[] v, float norm) { + float[] out = new float[dimension]; + for (int i = 0; i < dimension; i++) { + out[i] = v[i] / norm; + } + return out; + } + + /** Applies the D×D rotation matrix to the (normalized) input. */ + private float[] rotate(float[] v) { + float[][] rotation = getRotationMatrix(); + float[] out = new float[dimension]; + for (int i = 0; i < dimension; i++) { + double acc = 0.0; + float[] row = rotation[i]; + for (int j = 0; j < dimension; j++) { + acc += (double) row[j] * v[j]; + } + out[i] = (float) acc; + } + return out; + } + + /** Packs sign(v[i]) into ceil(D/8) bytes. Positive = bit 1, non-positive = bit 0. */ + private byte[] binarize(float[] v) { + byte[] code = new byte[(dimension + 7) / 8]; + for (int i = 0; i < dimension; i++) { + if (v[i] > 0f) { + code[i >> 3] |= (byte) (1 << (i & 7)); + } + } + return code; + } + + /** Lazy double-checked rotation matrix construction (Modified Gram-Schmidt). */ + private float[][] getRotationMatrix() { + if (rotMat == null) { + synchronized (this) { + if (rotMat == null) { + rotMat = getOrBuildRotationMatrix(dimension, seed); + } + } + } + return rotMat; + } + + private static float[][] getOrBuildRotationMatrix(int dimension, long seed) { + return ROTATION_MATRIX_CACHE.computeIfAbsent(new RotationKey(dimension, seed), + key -> buildRotationMatrix(key.dimension, key.seed)); + } + + /** + * Builds a random orthogonal D×D matrix using Modified Gram-Schmidt (MGS). + * Runtime: O(D³) — acceptable as a one-time cost (~0.5 s at D=768). + * Memory: 4 * D² bytes (~2.3 MB at D=768). + */ + static float[][] buildRotationMatrix(int d, long seed) { + Random rng = new Random(seed); + // Gaussian random matrix + double[][] randMatrix = new double[d][d]; + for (int i = 0; i < d; i++) { + for (int j = 0; j < d; j++) { + randMatrix[i][j] = rng.nextGaussian(); + } + } + // Modified Gram-Schmidt orthogonalization (column-wise) + for (int j = 0; j < d; j++) { + // Orthogonalize column j against all previous columns + for (int k = 0; k < j; k++) { + double dot = 0.0; + for (int i = 0; i < d; i++) { + dot += randMatrix[i][k] * randMatrix[i][j]; + } + for (int i = 0; i < d; i++) { + randMatrix[i][j] -= dot * randMatrix[i][k]; + } + } + // Normalize column j + double colNorm = 0.0; + for (int i = 0; i < d; i++) { + colNorm += randMatrix[i][j] * randMatrix[i][j]; + } + colNorm = Math.sqrt(colNorm); + if (colNorm > 1e-10) { + for (int i = 0; i < d; i++) { + randMatrix[i][j] /= colNorm; + } + } + } + // Transpose to row-major for efficient row-vector multiply + float[][] rotMat = new float[d][d]; + for (int i = 0; i < d; i++) { + for (int j = 0; j < d; j++) { + rotMat[i][j] = (float) randMatrix[j][i]; + } + } + return rotMat; + } + + private float[] subtract(float[] left, float[] right) { + float[] out = new float[dimension]; + for (int i = 0; i < dimension; i++) { + out[i] = left[i] - right[i]; + } + return out; + } + + private QuantizedLevels quantizeEx(float[] absValues, int exBits) { + double rescale = bestRescaleFactor(absValues, exBits); + int[] levels = new int[dimension]; + double ipNorm = 0.0; + int maxLevel = (1 << exBits) - 1; + for (int i = 0; i < dimension; i++) { + int level = (int) Math.floor((rescale * absValues[i]) + 1.0e-5); + if (level < 0) { + level = 0; + } else if (level > maxLevel) { + level = maxLevel; + } + levels[i] = level; + ipNorm += (level + 0.5d) * absValues[i]; + } + double ipNormInv = (Double.isFinite(ipNorm) && Math.abs(ipNorm) > 1.0e-12d) ? (1.0d / ipNorm) : 1.0d; + return new QuantizedLevels(levels, ipNormInv); + } + + private static float sum(float[] values) { + float sum = 0.0f; + for (float value : values) { + sum += value; + } + return sum; + } + + private double bestRescaleFactor(float[] absValues, int exBits) { + double max = 0.0d; + for (float value : absValues) { + max = Math.max(max, value); + } + if (max <= 1.0e-12d) { + return 1.0d; + } + + int maxLevel = (1 << exBits) - 1; + double bestFactor = maxLevel / max; + double bestScore = Double.NEGATIVE_INFINITY; + for (int step = 1; step <= 64; step++) { + double factor = (step * maxLevel) / (64.0d * max); + double numerator = 0.0d; + double denominator = dimension * 0.25d; + for (float value : absValues) { + int level = (int) Math.floor((factor * value) + 1.0e-5d); + if (level < 0) { + level = 0; + } else if (level > maxLevel) { + level = maxLevel; + } + numerator += (level + 0.5d) * value; + denominator += (level * (double) level) + level; + } + double score = numerator / Math.sqrt(denominator); + if (score > bestScore) { + bestScore = score; + bestFactor = factor; + } + } + return bestFactor; + } + + private byte[] packUnsignedLevels(int[] levels, int bitsPerValue) { + if (bitsPerValue <= 0) { + return new byte[0]; + } + byte[] packed = new byte[(dimension * bitsPerValue + 7) / 8]; + int bitOffset = 0; + for (int level : levels) { + for (int bit = 0; bit < bitsPerValue; bit++) { + if ((level & (1 << bit)) != 0) { + int absoluteBit = bitOffset + bit; + packed[absoluteBit >> 3] |= (byte) (1 << (absoluteBit & 7)); + } + } + bitOffset += bitsPerValue; + } + return packed; + } + + // ---- inner types ------------------------------------------------------- + + private static final class QuantizedLevels { + private final int[] levels; + private final double ipNormInv; + + private QuantizedLevels(int[] levels, double ipNormInv) { + this.levels = levels; + this.ipNormInv = ipNormInv; + } + } + + private static final class RotationKey { + private final int dimension; + private final long seed; + + private RotationKey(int dimension, long seed) { + this.dimension = dimension; + this.seed = seed; + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (other == null || getClass() != other.getClass()) { + return false; + } + RotationKey that = (RotationKey) other; + return dimension == that.dimension && seed == that.seed; + } + + @Override + public int hashCode() { + return Objects.hash(dimension, seed); + } + } +} \ No newline at end of file diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQFactorConfig.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQFactorConfig.java new file mode 100644 index 0000000000000..509eb2102f2be --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQFactorConfig.java @@ -0,0 +1,83 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +import java.io.Serializable; + +/** + * Generation-scoped configuration for RaBitQ neutral-factor computation (RFC-109 §3). + * + *

Replaces the previously static constants ({@code ERR_KAPPA}, absolute {@code EPS_IP}) with + * explicit, versioned thresholds. Readers reject factor versions they do not support instead of + * interpreting persisted factors with current-code constants. + * + *

+ */ +public final class RaBitQFactorConfig implements Serializable { + + private static final long serialVersionUID = 1L; + + public static final int FACTOR_VERSION = 1; + public static final double DEFAULT_KAPPA = 1.9; + public static final double DEFAULT_GMIN = 1.0e-3; + public static final double DEFAULT_EPS1_MAX = 1.0; + public static final double DEFAULT_EPS_N_REL = 1.0e-3; + + private final int factorVersion; + private final double kappa; + private final double gMin; + private final double eps1Max; + private final double epsNRel; + + public RaBitQFactorConfig(int factorVersion, double kappa, double gMin, double eps1Max, double epsNRel) { + if (factorVersion != FACTOR_VERSION) { + throw new IllegalArgumentException("Unsupported RaBitQ factor version: " + factorVersion); + } + this.factorVersion = factorVersion; + this.kappa = kappa; + this.gMin = gMin; + this.eps1Max = eps1Max; + this.epsNRel = epsNRel; + } + + /** The defaults for the first persisted factor format (RFC-109 §3). */ + public static RaBitQFactorConfig defaults() { + return new RaBitQFactorConfig( + FACTOR_VERSION, DEFAULT_KAPPA, DEFAULT_GMIN, DEFAULT_EPS1_MAX, DEFAULT_EPS_N_REL); + } + + public int getFactorVersion() { + return factorVersion; + } + + public double getKappa() { + return kappa; + } + + public double getGMin() { + return gMin; + } + + public double getEps1Max() { + return eps1Max; + } + + public double getEpsNRel() { + return epsNRel; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQNeutralFactors.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQNeutralFactors.java new file mode 100644 index 0000000000000..4faf8810d2815 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQNeutralFactors.java @@ -0,0 +1,133 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +/** + * Metric-NEUTRAL factor computation for residual RaBitQ encoding (RFC-109 §3). + * + *

The stored factors commit to NO metric; L2 / DOT / COSINE are composed at query time by + * {@link MetricQueryState}. The enclosing posting block's format version owns the factor layout: + * 

+ *   CENTER_RIP     = <c_rot, r>               exact, pass-independent
+ *   F_RESCALE_1    = n^2 / ipResidual1          pass-1 rescale
+ *   ERR_1          = n * eps1  (or n)           pass-1 absolute-IP error scale
+ *   F_RESCALE_EX   = n^2 / ipResidualEx         pass-2 rescale
+ *   RESIDUAL_NORM  = n
+ *   [ VECTOR_NORM ] = ||x||                     raw-cosine generations only
+ *
+ * + *

Corrections over the previous implementation: + *

+ */ +public final class RaBitQNeutralFactors { + + public static final double EPS_NORM = 1.0e-9; + + private RaBitQNeutralFactors() { + } + + /** Immutable per-vector factor set in posting-block scalar-array order. */ + public static final class Factors { + public final int factorVersion; + public final float centerRip; + public final float fRescale1; + public final float err1; + public final float fRescaleEx; + public final float residualNorm; + public final float vectorNorm; // consumed only for raw-cosine generations + + Factors(int factorVersion, float centerRip, float fRescale1, float err1, float fRescaleEx, + float residualNorm, float vectorNorm) { + this.factorVersion = factorVersion; + this.centerRip = centerRip; + this.fRescale1 = fRescale1; + this.err1 = err1; + this.fRescaleEx = fRescaleEx; + this.residualNorm = residualNorm; + this.vectorNorm = vectorNorm; + } + } + + /** + * @param residual r = rotated(x) - rotatedCenter, full precision (zero in padded dims) + * @param rotatedCenter P * c + * @param rotatedVector P * x (for ||x||; rotation preserves the norm) + * @param ipResidual1 <r, code1> accumulated in the encode loop + * @param ipResidualEx <r, codeEx> accumulated in the encode loop + * @param dimPadded padded dimension D'; padded coords are zero in r and codes + * @param config generation factor configuration (thresholds + version) + */ + public static Factors compute(float[] residual, float[] rotatedCenter, float[] rotatedVector, + double ipResidual1, double ipResidualEx, int dimPadded, + RaBitQFactorConfig config) { + int factorVersion = config.getFactorVersion(); + double nSq = 0.0; + double centerRip = 0.0; + double vSq = 0.0; + for (int i = 0; i < dimPadded; i++) { + nSq += (double) residual[i] * residual[i]; + centerRip += (double) rotatedCenter[i] * residual[i]; + vSq += (double) rotatedVector[i] * rotatedVector[i]; + } + double n = Math.sqrt(nSq); + float vectorNorm = (float) Math.sqrt(vSq); + // Residual-norm tiers (RFC-109 §3). + if (n == 0.0) { + // Vector coincides with centroid exactly: composition is exact; ERR_1 = 0 is legitimate. + return new Factors(factorVersion, (float) centerRip, 0f, 0f, 0f, 0f, vectorNorm); + } + if (n <= config.getEpsNRel() * vectorNorm) { + // Tiny but nonzero residual: disable the estimator, ERR_1 = residualNorm (maximal valid bound). + return new Factors(factorVersion, (float) centerRip, 0f, (float) n, 0f, (float) n, vectorNorm); + } + + // Normalized alignment of the residual with the sign code; ||code1|| = 0.5 * sqrt(dimPadded). + double gHat1 = ipResidual1 / (n * 0.5 * Math.sqrt(dimPadded)); + + float fRescale1; + float err1; + if (dimPadded == 1) { + // In one dimension the sign code reconstructs the residual exactly; the generic + // concentration bound is undefined because it contains (D - 1) in the denominator. + fRescale1 = (float) (nSq / ipResidual1); + err1 = 0f; + } else if (gHat1 < config.getGMin()) { + // Alignment too weak to trust the estimator (also guards the divide): disable, maximal bound. + fRescale1 = 0f; + err1 = (float) n; + } else { + double eps1 = config.getKappa() * Math.sqrt( + Math.max(0.0, 1.0 - gHat1 * gHat1) / (gHat1 * gHat1 * (dimPadded - 1))); + if (eps1 > config.getEps1Max()) { + // Relative error exceeds budget: disabling is the ONLY valid choice (never clamp to 1). + fRescale1 = 0f; + err1 = (float) n; + } else { + fRescale1 = (float) (nSq / ipResidual1); + err1 = (float) (n * eps1); + } + } + + // Pass-2 rescale; relative guard against a degenerate extended inner product (avoids Inf/NaN). + float fRescaleEx = Math.abs(ipResidualEx) <= EPS_NORM * Math.max(1.0, nSq) + ? 0f : (float) (nSq / ipResidualEx); + + return new Factors(factorVersion, (float) centerRip, fRescale1, err1, fRescaleEx, (float) n, vectorNorm); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQPlaneKernel.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQPlaneKernel.java new file mode 100644 index 0000000000000..c0b1f43d9a3b9 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQPlaneKernel.java @@ -0,0 +1,118 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +/** + * Direct plane popcount kernel (RFC-109 §3). Scores a quantized query ({@link VectorQueryPlanes}) + * against a data code stored as bit-planes (one sign plane + {@code exBits} extended planes, each a + * {@code long[]}), computing the inner product {@code } purely with + * {@code Long.bitCount(queryWord & dataWord)} — no per-dimension unpacking, sign-row copies, or + * extended-level repacking. + * + *

Decomposition. With {@code qPrime[i] = qMin + qLevel[i]*deltaQ} and, per dimension, + * {@code dataCode[i] = 2^exBits * signBit[i] + exLevel[i] + cBias}: + * 

+ *   <qPrime, dataCode> = qMin * sumDataCode + deltaQ * <qLevel, dataCode>
+ *
+ * where every {@code sum}/{@code <.,.>} term is a weighted popcount over plane words. The result is + * bit-exact with a scalar reference dot of the same {@code qPrime} against the same code. + */ +public final class RaBitQPlaneKernel { + + private RaBitQPlaneKernel() { + } + + /** + * Pass-1 inner product {@code } where {@code code1[i] = signBit[i] - 0.5} + * (sign-only 1-bit code). Uses only the sign plane. + */ + public static double scorePass1(VectorQueryPlanes q, long[] signPlane) { + int words = q.words(); + checkWords(signPlane, words, "signPlane"); + double qDotSignFull = queryDotPlane(q, signPlane); // + return qDotSignFull - 0.5 * q.sumQPrime(); + } + + /** + * Pass-2 inner product {@code } for the full {@code (exBits+1)}-bit centered + * code. {@code exPlanes[b]} is the b-th extended bit-plane (LSB first), each {@code long[words]}. + */ + public static double scorePass2(VectorQueryPlanes q, long[] signPlane, long[][] exPlanes, int exBits) { + int words = q.words(); + checkWords(signPlane, words, "signPlane"); + if (exBits <= 0) { + return scorePass1(q, signPlane); + } + if (exPlanes == null || exPlanes.length != exBits) { + throw new IllegalArgumentException("expected " + exBits + " extended planes"); + } + + double signScale = (double) (1L << exBits); + double cBias = -((double) (1L << (exBits + 1)) - 1.0) / 2.0; // -((2^bits - 1)/2), bits = exBits+1 + + // sumDataCode = 2^exBits * popcount(sign) + sum_b 2^b * popcount(exPlane_b) + cBias * dim + long popSign = popcount(signPlane); + double sumExLevels = 0.0; + for (int b = 0; b < exBits; b++) { + checkWords(exPlanes[b], words, "exPlane[" + b + "]"); + sumExLevels += (double) (1L << b) * popcount(exPlanes[b]); + } + double sumDataCode = signScale * popSign + sumExLevels + cBias * q.dim(); + + // = 2^exBits * + + cBias * sumQLevel + double qlDotSign = queryLevelDotPlane(q, signPlane); + double qlDotEx = 0.0; + for (int b = 0; b < exBits; b++) { + qlDotEx += (double) (1L << b) * queryLevelDotPlane(q, exPlanes[b]); + } + double sumQLevel = queryLevelSum(q); + double qlDotDataCode = signScale * qlDotSign + qlDotEx + cBias * sumQLevel; + + return q.queryMin() * sumDataCode + q.deltaQ() * qlDotDataCode; + } + + // = qMin*popcount(plane) + deltaQ * + private static double queryDotPlane(VectorQueryPlanes q, long[] plane) { + return q.queryMin() * popcount(plane) + q.deltaQ() * queryLevelDotPlane(q, plane); + } + + // = sum_a 2^a * popcount(queryPlane_a & plane) + private static double queryLevelDotPlane(VectorQueryPlanes q, long[] plane) { + double acc = 0.0; + for (int a = 0; a < q.bq(); a++) { + long[] qp = q.plane(a); + long pop = 0; + for (int w = 0; w < plane.length; w++) { + pop += Long.bitCount(qp[w] & plane[w]); + } + acc += (double) (1L << a) * pop; + } + return acc; + } + + // sum_i qLevel[i] = sum_a 2^a * popcount(queryPlane_a) + private static double queryLevelSum(VectorQueryPlanes q) { + double acc = 0.0; + for (int a = 0; a < q.bq(); a++) { + acc += (double) (1L << a) * popcount(q.plane(a)); + } + return acc; + } + + private static long popcount(long[] words) { + long pop = 0; + for (long w : words) { + pop += Long.bitCount(w); + } + return pop; + } + + private static void checkWords(long[] plane, int words, String name) { + if (plane == null || plane.length != words) { + throw new IllegalArgumentException(name + " must have " + words + " words"); + } + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQQueryState.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQQueryState.java new file mode 100644 index 0000000000000..0e3992a3d6489 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQQueryState.java @@ -0,0 +1,71 @@ +/* + * 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.hudi.common.index.vector; + +import java.io.Serializable; + +/** Query state reused while scoring RaBitQ candidates. */ +public final class RaBitQQueryState implements Serializable { + + private static final long serialVersionUID = 1L; + + private final byte[] binaryCode; + /** Rotation of the unit query, used by the asymmetric estimator. */ + private final float[] rotatedQuery; + /** Query norm before normalization, used for metric reconstruction. */ + private final float queryNorm; + /** Sum of {@link #rotatedQuery}, reused by multibit scoring. */ + private final float querySum; + + RaBitQQueryState(byte[] binaryCode, float[] rotatedQuery, float queryNorm) { + this.binaryCode = binaryCode.clone(); + this.rotatedQuery = rotatedQuery.clone(); + this.queryNorm = queryNorm; + float sum = 0f; + for (float value : rotatedQuery) { + sum += value; + } + this.querySum = sum; + } + + public byte[] getBinaryCode() { + return binaryCode.clone(); + } + + public float[] getRotatedQuery() { + return rotatedQuery.clone(); + } + + public float getQueryNorm() { + return queryNorm; + } + + public float getQuerySum() { + return querySum; + } + + byte[] binaryCodeUnsafe() { + return binaryCode; + } + + float[] rotatedQueryUnsafe() { + return rotatedQuery; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorDistanceMetric.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorDistanceMetric.java new file mode 100644 index 0000000000000..abbe51bf3c0f4 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorDistanceMetric.java @@ -0,0 +1,113 @@ +/* + * 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.hudi.common.index.vector; + +import java.util.Locale; + +/** + * Distance metrics for vector similarity search. + * + *

All metrics are returned as distances (smaller = more similar), + * so they can be compared uniformly with a min-heap. + */ +public enum VectorDistanceMetric { + + /** + * Cosine distance: 1 - cosine_similarity. + * Range: [0, 2]. 0 = identical direction, 2 = opposite. + */ + COSINE { + @Override + public float compute(float[] a, float[] b) { + checkDimensions(a, b); + double dot = 0; + double normA = 0; + double normB = 0; + for (int i = 0; i < a.length; i++) { + dot += (double) a[i] * b[i]; + normA += (double) a[i] * a[i]; + normB += (double) b[i] * b[i]; + } + double denom = Math.sqrt(normA) * Math.sqrt(normB); + return denom == 0.0 ? 1.0f : (float) (1.0 - dot / denom); + } + }, + + /** + * Euclidean (L2) distance. + * Range: [0, ∞). 0 = identical. + */ + L2 { + @Override + public float compute(float[] a, float[] b) { + checkDimensions(a, b); + double sum = 0; + for (int i = 0; i < a.length; i++) { + double d = (double) a[i] - b[i]; + sum += d * d; + } + return (float) Math.sqrt(sum); + } + }, + + /** + * Maximum inner product distance: negated dot product. + * Negated so smaller = higher similarity, consistent with the min-heap contract. + */ + DOT_PRODUCT { + @Override + public float compute(float[] a, float[] b) { + checkDimensions(a, b); + double dot = 0; + for (int i = 0; i < a.length; i++) { + dot += (double) a[i] * b[i]; + } + return (float) -dot; + } + }; + + /** + * Compute the distance between two float vectors. + * + * @param a first vector + * @param b second vector + * @return non-negative distance (smaller = more similar) + */ + public abstract float compute(float[] a, float[] b); + + /** + * Parse a metric name (case-insensitive). + * + * @param name e.g. "cosine", "l2", "dot_product" + * @return matching enum constant + */ + static VectorDistanceMetric fromString(String name) { + return valueOf(name.toUpperCase(Locale.ROOT).replace(" ", "_").replace("-", "_")); + } + + // ---- helpers ----------------------------------------------------------- + + private static void checkDimensions(float[] a, float[] b) { + if (a.length != b.length) { + throw new IllegalArgumentException( + "Vector dimension mismatch: " + a.length + " vs " + b.length); + } + } +} \ No newline at end of file diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexArbiter.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexArbiter.java new file mode 100644 index 0000000000000..a46449759234f --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexArbiter.java @@ -0,0 +1,130 @@ +/* + * 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.hudi.common.index.vector; + +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import java.util.Objects; + +/** + * The RLI finalist arbiter: resolves whether a vector-index posting still faithfully represents + * a live record, using the record-level index as the table's global version authority. + * + *

This is the pure classification core of RFC-109 "Upsert and Delete Support". It has no + * Spark or metadata-table dependency: callers resolve the current RLI location for a finalist + * key (a batched {@code readRecordIndexLocationsWithKeys} in the plan builder), then call + * {@link #classify} with the posting locator and that current location. The action taken per + * verdict is mode-specific and lives in the caller: + * + *

    + *
  • Approximate mode: {@code SERVE} -> keep, {@code STALE} -> exclude, {@code DELETED} -> exclude.
    • + *
  • Exact mode: {@code SERVE} -> positional fetch, {@code STALE} -> key-based fallback fetch, + * {@code DELETED} -> exclude.
    • + *
+ * + *

The classification itself is identical across modes so both report the same semantics and + * the same {@link ExclusionCounts} ({@code arbiterExclusions.stale} / {@code .deleted}). + */ +public final class VectorIndexArbiter { + + /** + * Verdict for a single finalist posting. + * + *

    + *
  • {@code SERVE}: RLI hit and the current location matches the posting locator. The posting + * faithfully represents the live record; positional trust is preserved (subject to the + * positional validity gate for {@code rowPosition >= 0}).
    • + *
  • {@code STALE}: RLI hit but the current location differs. The record was rewritten or + * moved; this posting's locator (and, for updates, its code) is no longer authoritative.
    • + *
  • {@code DELETED}: RLI miss. The record no longer exists in the table.
    • + *
+ */ + public enum Decision { + SERVE, + STALE, + DELETED + } + + private VectorIndexArbiter() { + } + + /** + * Classify a single finalist posting against the record's current RLI location. + * + * @param postingPartitionPath partition path stored in the posting locator (may be null) + * @param postingFileGroupId file group id stored in the posting locator (may be null) + * @param postingBaseInstantTime base instant time stored in the posting locator (may be null) + * @param currentLocation current RLI location for the record key, or {@code null} for an + * RLI miss (deleted) + * @return the arbiter verdict + */ + public static Decision classify(String postingPartitionPath, + String postingFileGroupId, + String postingBaseInstantTime, + HoodieRecordGlobalLocation currentLocation) { + if (currentLocation == null) { + return Decision.DELETED; + } + boolean matches = Objects.equals(postingPartitionPath, currentLocation.getPartitionPath()) + && Objects.equals(postingFileGroupId, currentLocation.getFileId()) + && Objects.equals(postingBaseInstantTime, currentLocation.getInstantTime()); + return matches ? Decision.SERVE : Decision.STALE; + } + + /** + * Mutable tally of arbiter exclusions, mirrored into both query modes' log lines as the + * {@code arbiterExclusions} freshness observability metric. Split into {@code stale} and + * {@code deleted} so the two upsert/delete effects are separable in dashboards. + */ + public static final class ExclusionCounts { + private long stale; + private long deleted; + + /** Record a verdict, incrementing the matching counter. {@code SERVE} is a no-op. */ + public void record(Decision decision) { + switch (decision) { + case STALE: + stale++; + break; + case DELETED: + deleted++; + break; + default: + break; + } + } + + public long stale() { + return stale; + } + + public long deleted() { + return deleted; + } + + public long total() { + return stale + deleted; + } + + @Override + public String toString() { + return "arbiterExclusions{stale=" + stale + ", deleted=" + deleted + "}"; + } + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexOptions.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexOptions.java new file mode 100644 index 0000000000000..44fd48cb5539b --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexOptions.java @@ -0,0 +1,243 @@ +/* + * 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.hudi.common.index.vector; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.Locale; +import java.util.Map; +import java.util.Set; + +/** + * Options accepted by {@code CREATE INDEX ... USING VECTOR}. + * + *

The indexed column's Hudi {@code VECTOR(D[, elementType])} schema is authoritative for + * dimension and element type. Index options configure only the acceleration structure. DDL + * implementations must call {@link #validateAndNormalize(Map)} before persisting an index + * definition; individual parsing helpers are intentionally private so aggregate validation cannot + * be bypassed. + */ +public final class VectorIndexOptions { + + public static final String METRIC = "vector.metric"; + public static final String QUANTIZER = "vector.quantizer"; + public static final String NUM_CLUSTERS = "vector.num_clusters"; + public static final String MAX_ITER = "vector.max_iter"; + public static final String RABITQ_BITS = "vector.rabitq.bits"; + public static final String RABITQ_SEED = "vector.rabitq.seed"; + public static final String RABITQ_ASSUME_NORMALIZED = "vector.rabitq.assume_normalized"; + public static final String QUERY_NUM_PROBES = "vector.query.nprobes"; + public static final String QUERY_REFINE_FACTOR = "vector.query.refine_factor"; + public static final String QUERY_MODE = "vector.query.mode"; + public static final String QUERY_STALE_POLICY = "vector.query.stale_policy"; + + public static final VectorDistanceMetric DEFAULT_METRIC = VectorDistanceMetric.COSINE; + public static final VectorQuantizer DEFAULT_QUANTIZER = VectorQuantizer.IVF_RABITQ; + public static final int DEFAULT_NUM_CLUSTERS = 256; + public static final int DEFAULT_MAX_ITER = 20; + public static final int DEFAULT_RABITQ_BITS = 4; + public static final long DEFAULT_RABITQ_SEED = 42L; + public static final int DEFAULT_NUM_PROBES = 32; + public static final int DEFAULT_REFINE_FACTOR = 50; + public static final VectorQueryMode DEFAULT_QUERY_MODE = VectorQueryMode.EXACT_RERANK; + public static final VectorStalePolicy DEFAULT_STALE_POLICY = VectorStalePolicy.FAIL; + + private static final Set SUPPORTED_OPTIONS = Collections.unmodifiableSet( + new HashSet<>(Arrays.asList( + METRIC, + QUANTIZER, + NUM_CLUSTERS, + MAX_ITER, + RABITQ_BITS, + RABITQ_SEED, + RABITQ_ASSUME_NORMALIZED, + QUERY_NUM_PROBES, + QUERY_REFINE_FACTOR, + QUERY_MODE, + QUERY_STALE_POLICY))); + + private VectorIndexOptions() { + } + + /** + * Validates the complete option map and returns canonical values for persistence. + * + *

The returned map contains every supported option, including explicit defaults. Unknown, + * retired, misspelled, and invalid options are rejected instead of being silently ignored. + */ + public static Map validateAndNormalize(Map options) { + Set unknownOptions = new HashSet<>(options.keySet()); + unknownOptions.removeAll(SUPPORTED_OPTIONS); + if (!unknownOptions.isEmpty()) { + throw new IllegalArgumentException("Unsupported vector index options: " + unknownOptions); + } + + VectorDistanceMetric metric = getMetric(options); + VectorQuantizer quantizer = getQuantizer(options); + int numClusters = getNumClusters(options); + int maxIter = getMaxIter(options); + int bits = getRaBitQBits(options); + long seed = getRaBitQSeed(options); + boolean assumeNormalized = shouldAssumeNormalizedVectors(options); + int numProbes = getNumProbes(options); + int refineFactor = getRefineFactor(options); + VectorQueryMode queryMode = getQueryMode(options); + VectorStalePolicy stalePolicy = getStalePolicy(options); + + if (numProbes > numClusters) { + throw new IllegalArgumentException( + "Option '" + QUERY_NUM_PROBES + "' must not exceed '" + NUM_CLUSTERS + "': " + + numProbes + " > " + numClusters); + } + + Map normalized = new LinkedHashMap<>(); + normalized.put(METRIC, metric.name().toLowerCase(Locale.ROOT)); + normalized.put(QUANTIZER, quantizer.name()); + normalized.put(NUM_CLUSTERS, String.valueOf(numClusters)); + normalized.put(MAX_ITER, String.valueOf(maxIter)); + normalized.put(RABITQ_BITS, String.valueOf(bits)); + normalized.put(RABITQ_SEED, String.valueOf(seed)); + normalized.put(RABITQ_ASSUME_NORMALIZED, String.valueOf(assumeNormalized)); + normalized.put(QUERY_NUM_PROBES, String.valueOf(numProbes)); + normalized.put(QUERY_REFINE_FACTOR, String.valueOf(refineFactor)); + normalized.put(QUERY_MODE, queryMode.name().toLowerCase(Locale.ROOT)); + normalized.put(QUERY_STALE_POLICY, stalePolicy.name().toLowerCase(Locale.ROOT)); + return Collections.unmodifiableMap(normalized); + } + + private static VectorDistanceMetric getMetric(Map options) { + String value = getOption(options, METRIC, DEFAULT_METRIC.name()); + try { + return VectorDistanceMetric.fromString(value); + } catch (IllegalArgumentException e) { + throw unsupportedValue(METRIC, value, e); + } + } + + private static VectorQuantizer getQuantizer(Map options) { + String value = getOption(options, QUANTIZER, DEFAULT_QUANTIZER.name()); + try { + return VectorQuantizer.fromString(value); + } catch (IllegalArgumentException e) { + throw unsupportedValue(QUANTIZER, value, e); + } + } + + private static int getNumClusters(Map options) { + return getPositiveInt(options, NUM_CLUSTERS, DEFAULT_NUM_CLUSTERS); + } + + private static int getMaxIter(Map options) { + return getPositiveInt(options, MAX_ITER, DEFAULT_MAX_ITER); + } + + private static int getRaBitQBits(Map options) { + int bits = getInt(options, RABITQ_BITS, DEFAULT_RABITQ_BITS); + if (bits < 1 || bits > 8) { + throw new IllegalArgumentException( + "Option '" + RABITQ_BITS + "' must be between 1 and 8: " + bits); + } + return bits; + } + + private static long getRaBitQSeed(Map options) { + String value = getOption(options, RABITQ_SEED, String.valueOf(DEFAULT_RABITQ_SEED)); + try { + return Long.parseLong(value); + } catch (NumberFormatException e) { + throw invalidNumber(RABITQ_SEED, value, e); + } + } + + private static boolean shouldAssumeNormalizedVectors(Map options) { + String value = getOption( + options, RABITQ_ASSUME_NORMALIZED, String.valueOf(false)).toLowerCase(Locale.ROOT); + if (!"true".equals(value) && !"false".equals(value)) { + throw new IllegalArgumentException( + "Option '" + RABITQ_ASSUME_NORMALIZED + "' must be either 'true' or 'false': " + value); + } + return Boolean.parseBoolean(value); + } + + private static int getNumProbes(Map options) { + return getPositiveInt(options, QUERY_NUM_PROBES, DEFAULT_NUM_PROBES); + } + + private static int getRefineFactor(Map options) { + return getPositiveInt(options, QUERY_REFINE_FACTOR, DEFAULT_REFINE_FACTOR); + } + + private static VectorQueryMode getQueryMode(Map options) { + String value = getOption(options, QUERY_MODE, DEFAULT_QUERY_MODE.name()); + try { + return VectorQueryMode.fromString(value); + } catch (IllegalArgumentException e) { + throw unsupportedValue(QUERY_MODE, value, e); + } + } + + private static VectorStalePolicy getStalePolicy(Map options) { + String value = getOption(options, QUERY_STALE_POLICY, DEFAULT_STALE_POLICY.name()); + try { + return VectorStalePolicy.fromString(value); + } catch (IllegalArgumentException e) { + throw unsupportedValue(QUERY_STALE_POLICY, value, e); + } + } + + private static int getPositiveInt(Map options, String key, int defaultValue) { + int value = getInt(options, key, defaultValue); + if (value <= 0) { + throw new IllegalArgumentException("Option '" + key + "' must be greater than 0: " + value); + } + return value; + } + + private static int getInt(Map options, String key, int defaultValue) { + String value = getOption(options, key, String.valueOf(defaultValue)); + try { + return Integer.parseInt(value); + } catch (NumberFormatException e) { + throw invalidNumber(key, value, e); + } + } + + private static IllegalArgumentException invalidNumber( + String key, String value, NumberFormatException cause) { + return new IllegalArgumentException( + "Option '" + key + "' must be a valid number: " + value, cause); + } + + private static IllegalArgumentException unsupportedValue( + String key, String value, IllegalArgumentException cause) { + return new IllegalArgumentException( + "Unsupported value for option '" + key + "': " + value, cause); + } + + private static String getOption(Map options, String key, String defaultValue) { + String value = options.getOrDefault(key, defaultValue); + if (value == null || value.isEmpty()) { + throw new IllegalArgumentException("Option '" + key + "' must not be empty"); + } + return value; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexPruner.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexPruner.java new file mode 100644 index 0000000000000..97a640cc216b8 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexPruner.java @@ -0,0 +1,154 @@ +/* + * 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.hudi.common.index.vector; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Engine-agnostic IVF cluster pruner for vector queries. + * + *

Given a set of centroids and a file-group-to-cluster mapping, determines which + * file groups need to be scanned for an approximate nearest-neighbour query. + * + *

Used by both the Spark file index ({@code HoodieVectorAwareFileIndex}) and the + * Trino split manager ({@code HudiVectorSplitManager}), keeping all math in one place. + * + *

Instances are built from MDT data at query planning time and are short-lived + * (one per query or per query batch). They are not cached between queries because + * centroids can change after LIRE compaction. + */ +public final class VectorIndexPruner implements Serializable { + + private static final long serialVersionUID = 1L; + + /** Centroid vectors keyed by cluster id (0-based). */ + private final float[][] centroids; + + /** + * Mapping from cluster id → set of file group ids containing vectors in that cluster. + * The inner sets are unmodifiable. + */ + private final Map> clusterToFileGroups; + + /** Distance metric for centroid scoring. */ + private final VectorDistanceMetric metric; + + /** + * @param centroids centroid vectors, indexed by cluster id + * @param clusterToFileGroups mapping built from the MDT fg_mapping partition + * @param metric distance metric matching the index definition + */ + public VectorIndexPruner( + float[][] centroids, + Map> clusterToFileGroups, + VectorDistanceMetric metric) { + this.centroids = centroids; + this.clusterToFileGroups = clusterToFileGroups; + this.metric = metric; + } + + /** + * Returns the set of file group ids that must be scanned to answer an ANN query. + * + * @param queryVector the query embedding + * @param numProbes number of clusters to probe (nProbes) + * @return file group ids; never null, may be empty if the index is not yet initialized + */ + public Set probe(float[] queryVector, int numProbes) { + if (centroids == null || centroids.length == 0) { + return Collections.emptySet(); + } + int effectiveProbes = Math.min(numProbes, centroids.length); + int[] topClusters = findTopClusters(queryVector, effectiveProbes); + + Set fileGroups = new HashSet<>(); + for (int clusterId : topClusters) { + Set fgs = clusterToFileGroups.get(clusterId); + if (fgs != null) { + fileGroups.addAll(fgs); + } + } + return Collections.unmodifiableSet(fileGroups); + } + + /** + * Returns the cluster ids closest to the query (sorted best-first). + * Linear scan; HNSW routing replaces this in Phase 2. + */ + public int[] findTopClusters(float[] query, int numProbes) { + int k = centroids.length; + // scored[i] = (distance, cluster_id) + List scored = new ArrayList<>(k); + for (int i = 0; i < k; i++) { + float dist = metric.compute(query, centroids[i]); + scored.add(new float[]{dist, i}); + } + scored.sort((a, b) -> Float.compare(a[0], b[0])); + + int[] result = new int[numProbes]; + for (int i = 0; i < numProbes; i++) { + result[i] = (int) scored.get(i)[1]; + } + return result; + } + + /** Returns the number of clusters (K). */ + public int numClusters() { + return centroids == null ? 0 : centroids.length; + } + + // ---- factory helpers --------------------------------------------------- + + /** + * Builds a cluster→file-group map from a flat assignment list. + * Each entry in {@code assignments} must be a three-element array: + * {@code [clusterId (int), fileGroupId (String), partitionPath (String)]}. + * + *

The returned map is partition-aware: pass {@code partitionFilter = null} + * to include all partitions, or a non-null set to restrict to specific partitions. + * + * @param assignments rows from the MDT assignments/fg_mapping partition + * @param partitionFilter partition paths to include; null = all + * @return cluster → file-group mapping + */ + public static Map> buildClusterMap( + Iterable assignments, + Set partitionFilter) { + Map> map = new HashMap<>(); + for (Object[] row : assignments) { + int clusterId = ((Number) row[0]).intValue(); + String fileGroupId = (String) row[1]; + String partitionPath = (String) row[2]; + + if (partitionFilter != null && !partitionFilter.contains(partitionPath)) { + continue; + } + map.computeIfAbsent(clusterId, k -> new HashSet<>()).add(fileGroupId); + } + return map; + } +} \ No newline at end of file diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorQuantizer.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorQuantizer.java new file mode 100644 index 0000000000000..56190d425f93f --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorQuantizer.java @@ -0,0 +1,31 @@ +/* + * 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.hudi.common.index.vector; + +import java.util.Locale; + +/** Supported vector-index quantizers. */ +public enum VectorQuantizer { + IVF_RABITQ; + + static VectorQuantizer fromString(String value) { + return valueOf(value.toUpperCase(Locale.ROOT).replace('-', '_')); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorQueryMode.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorQueryMode.java new file mode 100644 index 0000000000000..f7601ab40ab38 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorQueryMode.java @@ -0,0 +1,32 @@ +/* + * 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.hudi.common.index.vector; + +import java.util.Locale; + +/** Supported vector-query execution modes. */ +public enum VectorQueryMode { + APPROXIMATE, + EXACT_RERANK; + + static VectorQueryMode fromString(String value) { + return valueOf(value.toUpperCase(Locale.ROOT).replace('-', '_')); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorQueryPlanes.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorQueryPlanes.java new file mode 100644 index 0000000000000..389f4e6ecd017 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorQueryPlanes.java @@ -0,0 +1,167 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +import java.io.Serializable; + +/** + * The probed-cluster query quantized to {@code Bq} bit-planes for direct popcount scoring + * (RFC-109 §3). The selected cluster query {@code wRot} is uniformly scalar-quantized to + * {@code Bq}-bit levels once per probed cluster; each bit position is stored as a {@code long[]} + * plane so the scan kernel can score with {@code Long.bitCount(dataWord & queryWord)} instead of + * per-dimension float multiply-add. + * + *

Reconstruction: {@code qPrime[i] = qMin + level[i] * deltaQ}. The scan kernel computes + * {@code } exactly; the residual {@code selectedQuery - qPrime} is captured by + * {@link #quantizationErrorNorm} and folds into the plane error bound + * {@code errAbs = ERR_1 * ||qPrime|| + RESIDUAL_NORM * ||selectedQuery - qPrime||}. + */ +public final class VectorQueryPlanes implements Serializable { + + private static final long serialVersionUID = 1L; + + private final long[][] planes; // [bq][words] + private final int dim; + private final int bq; + private final int words; + private final double qMin; + private final double deltaQ; + private final double sumQPrime; + private final double qPrimeNorm; + private final double quantizationErrorNorm; + + private VectorQueryPlanes(long[][] planes, int dim, int bq, int words, double qMin, double deltaQ, + double sumQPrime, double qPrimeNorm, double quantizationErrorNorm) { + this.planes = planes; + this.dim = dim; + this.bq = bq; + this.words = words; + this.qMin = qMin; + this.deltaQ = deltaQ; + this.sumQPrime = sumQPrime; + this.qPrimeNorm = qPrimeNorm; + this.quantizationErrorNorm = quantizationErrorNorm; + } + + /** + * Quantize the (rotated) selected cluster query to {@code bq} bit-planes. + * + * @param selectedQuery the rotated residual query {@code wRot} (length = dimPadded) + * @param bq query bit width (RFC-109 default {@code Bq = 4}) + */ + public static VectorQueryPlanes quantize(float[] selectedQuery, int bq) { + if (bq <= 0 || bq > 16) { + throw new IllegalArgumentException("Bq must be in [1, 16], got: " + bq); + } + int dim = selectedQuery.length; + int words = (dim + 63) >> 6; + int maxLevel = (1 << bq) - 1; + + double min = Double.POSITIVE_INFINITY; + double max = Double.NEGATIVE_INFINITY; + for (float v : selectedQuery) { + if (v < min) { + min = v; + } + if (v > max) { + max = v; + } + } + double deltaQ = (max > min) ? (max - min) / maxLevel : 0.0; + + long[][] planes = new long[bq][words]; + double sumQPrime = 0.0; + double qPrimeNormSq = 0.0; + double errNormSq = 0.0; + for (int i = 0; i < dim; i++) { + int level = 0; + if (deltaQ > 0.0) { + level = (int) Math.round((selectedQuery[i] - min) / deltaQ); + if (level < 0) { + level = 0; + } else if (level > maxLevel) { + level = maxLevel; + } + } + double qp = min + level * deltaQ; + int w = i >> 6; + long bit = 1L << (i & 63); + for (int b = 0; b < bq; b++) { + if (((level >> b) & 1) != 0) { + planes[b][w] |= bit; + } + } + sumQPrime += qp; + qPrimeNormSq += qp * qp; + double d = selectedQuery[i] - qp; + errNormSq += d * d; + } + + return new VectorQueryPlanes(planes, dim, bq, words, min, deltaQ, sumQPrime, + Math.sqrt(qPrimeNormSq), Math.sqrt(errNormSq)); + } + + /** Reconstruct the quantized query {@code qPrime} (mainly for tests / scalar reference). */ + public float[] reconstruct() { + float[] out = new float[dim]; + for (int i = 0; i < dim; i++) { + int w = i >> 6; + long bit = 1L << (i & 63); + int level = 0; + for (int b = 0; b < bq; b++) { + if ((planes[b][w] & bit) != 0) { + level |= (1 << b); + } + } + out[i] = (float) (qMin + level * deltaQ); + } + return out; + } + + public long[] plane(int b) { + return planes[b]; + } + + public int dim() { + return dim; + } + + public int bq() { + return bq; + } + + public int words() { + return words; + } + + public double queryMin() { + return qMin; + } + + public double deltaQ() { + return deltaQ; + } + + public double sumQPrime() { + return sumQPrime; + } + + public double queryPrimeNorm() { + return qPrimeNorm; + } + + public double quantizationErrorNorm() { + return quantizationErrorNorm; + } + + /** + * Plane error bound (RFC-109 §3): + * {@code errAbs = err1 * ||qPrime|| + residualNorm * ||selectedQuery - qPrime||}. + */ + public double planeErrorBound(double err1, double residualNorm) { + return err1 * qPrimeNorm + residualNorm * quantizationErrorNorm; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorStalePolicy.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorStalePolicy.java new file mode 100644 index 0000000000000..75a225923d9c2 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorStalePolicy.java @@ -0,0 +1,33 @@ +/* + * 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.hudi.common.index.vector; + +import java.util.Locale; + +/** Policies applied when vector-index state is stale at the pinned snapshot. */ +public enum VectorStalePolicy { + FAIL, + WARN, + FALLBACK; + + static VectorStalePolicy fromString(String value) { + return valueOf(value.toUpperCase(Locale.ROOT)); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ArbitratedVectorCandidate.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ArbitratedVectorCandidate.java new file mode 100644 index 0000000000000..63e37bad4230d --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ArbitratedVectorCandidate.java @@ -0,0 +1,69 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import java.io.Serializable; + +/** + * A candidate after RLI freshness arbitration (RFC-109 §7): the original candidate, its + * {@link VectorCandidateState} verdict, and the live location resolved from the RLI (present for + * {@code SERVE} and {@code STALE}, null for {@code DELETED}). + * + *

    + *
  • {@code SERVE}: {@code liveLocation} matches the posting hint; positional read allowed.
    • + *
  • {@code STALE}: {@code liveLocation} is the record's current location; exact mode must + * key-fetch there rather than trust the posting's row position.
    • + *
  • {@code DELETED}: dropped; no live location.
    • + *
+ */ +public final class ArbitratedVectorCandidate implements Serializable { + + private static final long serialVersionUID = 1L; + + private final VectorCandidate candidate; + private final VectorCandidateState state; + private final HoodieRecordGlobalLocation liveLocation; + + public ArbitratedVectorCandidate(VectorCandidate candidate, + VectorCandidateState state, + HoodieRecordGlobalLocation liveLocation) { + this.candidate = candidate; + this.state = state; + this.liveLocation = liveLocation; + } + + public VectorCandidate getCandidate() { + return candidate; + } + + public VectorCandidateState getState() { + return state; + } + + /** Live RLI location for SERVE/STALE, or null for DELETED. */ + public HoodieRecordGlobalLocation getLiveLocation() { + return liveLocation; + } + + public boolean isServable() { + return state == VectorCandidateState.SERVE || state == VectorCandidateState.STALE; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/CommonVectorSearchExecutor.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/CommonVectorSearchExecutor.java new file mode 100644 index 0000000000000..ea24a85e3c2f0 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/CommonVectorSearchExecutor.java @@ -0,0 +1,92 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.data.HoodieListData; +import org.apache.hudi.common.engine.HoodieEngineContext; + +import java.util.Objects; + +/** + * The engine-neutral vector-search orchestrator (RFC-109 §11). Pins one snapshot, chooses the + * execution mode, and drives the stages in order — candidate scan, RLI arbitration, fetch planning, + * exact rerank — returning the top-K results. It never invokes {@code spark.sql(...)} and never + * reconstructs SQL/DataFrames; each stage is an injected engine-neutral implementation. + * + * 
+ *   request
+ *     -> resolve + pin snapshot
+ *     -> select execution mode (LOCAL/DISTRIBUTED)
+ *     -> scan MDT postings (bounded ordered pool)
+ *     -> RLI arbitrate (SERVE/STALE, drop DELETED)
+ *     -> plan file-slice fetches
+ *     -> exact rerank + continuation
+ *     -> top-K results
+ *
+ */ +public final class CommonVectorSearchExecutor implements VectorSearchExecutor { + + private static final long serialVersionUID = 1L; + + private final VectorSnapshotResolver snapshotResolver; + private final VectorExecutionModeSelector executionModeSelector; + private final VectorCandidateSource candidateSource; + private final VectorCandidateArbiter candidateArbiter; + private final VectorFetchPlanner fetchPlanner; + private final VectorExactReranker exactReranker; + + public CommonVectorSearchExecutor(VectorSnapshotResolver snapshotResolver, + VectorExecutionModeSelector executionModeSelector, + VectorCandidateSource candidateSource, + VectorCandidateArbiter candidateArbiter, + VectorFetchPlanner fetchPlanner, + VectorExactReranker exactReranker) { + this.snapshotResolver = Objects.requireNonNull(snapshotResolver, "snapshotResolver"); + this.executionModeSelector = Objects.requireNonNull(executionModeSelector, "executionModeSelector"); + this.candidateSource = Objects.requireNonNull(candidateSource, "candidateSource"); + this.candidateArbiter = Objects.requireNonNull(candidateArbiter, "candidateArbiter"); + this.fetchPlanner = Objects.requireNonNull(fetchPlanner, "fetchPlanner"); + this.exactReranker = Objects.requireNonNull(exactReranker, "exactReranker"); + } + + @Override + public HoodieData execute(VectorSearchRequest request, HoodieEngineContext engineContext) { + // 1. Pin one snapshot for MDT / RLI / file-slice / base reads. + VectorSearchSnapshot snapshot = snapshotResolver.resolve(request); + // 2. Choose execution locality (recorded on the plan for downstream + metrics). + VectorExecutionDecision decision = executionModeSelector.select(request); + VectorSearchPlan plan = new VectorSearchPlan(request, snapshot, decision); + + // 3-6. Consume ordered windows from one retained pool. Continuation never rescans MDT. + VectorCandidatePool pool = candidateSource.scan(plan, engineContext); + VectorTopKAccumulator results = new VectorTopKAccumulator(request.getTopK()); + while (results.needsMore() && pool.hasMore()) { + HoodieData candidates = pool.nextBatch(); + HoodieData arbitrated = + candidateArbiter.arbitrate(candidates, request, snapshot, engineContext); + HoodieData tasks = fetchPlanner.plan(arbitrated, snapshot, engineContext); + for (VectorSearchResult result : exactReranker.rerank( + tasks, request, snapshot, engineContext).collectAsList()) { + results.offer(result.getRecordKey(), result.getDistance(), result.getLocation()); + } + } + return HoodieListData.eager(results.topK()); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DeadlinePolicy.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DeadlinePolicy.java new file mode 100644 index 0000000000000..ccb88c79b2346 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DeadlinePolicy.java @@ -0,0 +1,30 @@ +/* + * 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.hudi.common.index.vector.search; + +/** + * What to do when a vector search cannot assemble K live exact results before the request deadline. + * {@code FAIL} surfaces a deadline error; {@code RETURN_PARTIAL} returns the live results gathered + * so far with a {@link VectorSearchStatus#DEADLINE_EXCEEDED} / {@link VectorSearchStatus#PARTIAL} + * status. Returning fewer than K silently is never allowed (RFC-109 §10). + */ +public enum DeadlinePolicy { + FAIL, + RETURN_PARTIAL +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DefaultExactVectorScorer.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DefaultExactVectorScorer.java new file mode 100644 index 0000000000000..fe3dd70262d40 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DefaultExactVectorScorer.java @@ -0,0 +1,85 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.index.vector.VectorDistanceMetric; + +/** + * Default exact scorer (RFC-109 §10). Accumulates in float64 and returns an order-preserving + * ranking distance (smaller = more similar), consistent with the approximate path: + * + *
    + *
  • {@code L2}: squared L2 (kept squared internally through ranking; callers may take the + * square root only for final presentation).
    • + *
  • {@code DOT_PRODUCT}: negated dot product.
    • + *
  • {@code COSINE}: {@code 1 - cosine_similarity}.
    • + *
+ */ +public final class DefaultExactVectorScorer implements ExactVectorScorer { + + private static final long serialVersionUID = 1L; + + @Override + public double distance(float[] query, float[] candidate, VectorDistanceMetric metric) { + if (query.length != candidate.length) { + throw new IllegalArgumentException( + "Vector dimension mismatch: " + query.length + " vs " + candidate.length); + } + switch (metric) { + case L2: + return squaredL2(query, candidate); + case DOT_PRODUCT: + return -dot(query, candidate); + case COSINE: + return cosineDistance(query, candidate); + default: + throw new IllegalArgumentException("Unsupported metric: " + metric); + } + } + + private static double squaredL2(float[] a, float[] b) { + double sum = 0.0; + for (int i = 0; i < a.length; i++) { + double d = (double) a[i] - (double) b[i]; + sum += d * d; + } + return sum; + } + + private static double dot(float[] a, float[] b) { + double dot = 0.0; + for (int i = 0; i < a.length; i++) { + dot += (double) a[i] * (double) b[i]; + } + return dot; + } + + private static double cosineDistance(float[] a, float[] b) { + double dot = 0.0; + double normA = 0.0; + double normB = 0.0; + for (int i = 0; i < a.length; i++) { + dot += (double) a[i] * (double) b[i]; + normA += (double) a[i] * (double) a[i]; + normB += (double) b[i] * (double) b[i]; + } + double denom = Math.sqrt(normA) * Math.sqrt(normB); + return denom == 0.0 ? 1.0 : 1.0 - dot / denom; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DefaultVectorExactReranker.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DefaultVectorExactReranker.java new file mode 100644 index 0000000000000..13be2bced58cb --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DefaultVectorExactReranker.java @@ -0,0 +1,113 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.data.HoodieListData; +import org.apache.hudi.common.engine.HoodieEngineContext; +import org.apache.hudi.common.index.vector.VectorDistanceMetric; + +import java.util.Iterator; +import java.util.List; + +/** + * Default engine-neutral exact reranker (RFC-109 §10). For each fetch task it reads the + * candidate rows through an injected {@link HoodieVectorBatchReadHandle} (created per partition via + * {@link HoodieVectorBatchReadHandleSupplier}), scores them with an {@link ExactVectorScorer} + * (float64, squared-L2 internal), and keeps a per-partition top-K via {@link VectorTopKAccumulator}. + * The bounded per-partition top-Ks are then merged into the global top-K — the only driver-side + * collect, and only of at most {@code partitions * topK} results (never the full candidate set), + * satisfying "collect only final top-K" for distributed execution. + * + *

This scores exactly the rows the {@link VectorFetchPlanner} produced (DELETED already excluded, + * STALE via key fallback). Pool-level continuation — re-drawing more candidates when stale/deleted + * crowding leaves fewer than K live — is orchestrated upstream in the executor using + * {@link VectorContinuationController}; this reranker scores a given batch of tasks. + */ +public final class DefaultVectorExactReranker implements VectorExactReranker { + + private static final long serialVersionUID = 1L; + + private final HoodieVectorBatchReadHandleSupplier handleSupplier; + private final ExactVectorScorer scorer; + private final String recordKeyField; + private final String vectorColumn; + + public DefaultVectorExactReranker(HoodieVectorBatchReadHandleSupplier handleSupplier, + ExactVectorScorer scorer, + String recordKeyField, + String vectorColumn) { + this.handleSupplier = handleSupplier; + this.scorer = scorer; + this.recordKeyField = recordKeyField; + this.vectorColumn = vectorColumn; + } + + @Override + public HoodieData rerank(HoodieData tasks, + VectorSearchRequest request, + VectorSearchSnapshot snapshot, + HoodieEngineContext engineContext) { + int topK = request.getTopK(); + float[] query = request.getQueryVector(); + VectorDistanceMetric metric = request.getMetric(); + HoodieVectorBatchReadHandleSupplier supplier = this.handleSupplier; + ExactVectorScorer localScorer = this.scorer; + String keyField = this.recordKeyField; + String vecCol = this.vectorColumn; + + // Per-partition: read + score + local top-K. + HoodieData perPartition = tasks.mapPartitions(taskIt -> { + VectorTopKAccumulator acc = new VectorTopKAccumulator(topK); + if (taskIt.hasNext()) { + HoodieVectorBatchReadHandle handle = supplier.get(); + try { + while (taskIt.hasNext()) { + VectorFetchTask task = taskIt.next(); + Iterator records = handle.read(task, keyField, vecCol); + while (records.hasNext()) { + VectorRecord rec = records.next(); + double dist = localScorer.distance(query, rec.getVector(), metric); + acc.offer(rec.getRecordKey(), dist, rec.getLocation()); + } + } + } finally { + closeQuietly(handle); + } + } + return acc.topK().iterator(); + }, false); + + // Global merge of the bounded per-partition top-Ks (final top-K only). + List merged = perPartition.collectAsList(); + VectorTopKAccumulator global = new VectorTopKAccumulator(topK); + for (VectorSearchResult r : merged) { + global.offer(r.getRecordKey(), r.getDistance(), r.getLocation()); + } + return HoodieListData.eager(global.topK()); + } + + private static void closeQuietly(HoodieVectorBatchReadHandle handle) { + try { + handle.close(); + } catch (Exception ignored) { + // best-effort close + } + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DefaultVectorFetchPlanner.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DefaultVectorFetchPlanner.java new file mode 100644 index 0000000000000..d1d2503a5288a --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/DefaultVectorFetchPlanner.java @@ -0,0 +1,92 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.engine.HoodieEngineContext; +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; +import org.apache.hudi.common.util.collection.Pair; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; + +/** + * Default engine-neutral fetch planner (RFC-109 §8). Groups arbitrated candidates by their live + * file (partition + fileId, resolved by the arbiter against the pinned snapshot) into one + * {@link VectorFetchTask} per file, so the read handle can coalesce positions within a file. + * + *

    + *
  • {@code SERVE}: the posting row position is preserved for a positional read.
    • + *
  • {@code STALE}: the row position is dropped ({@code -1}); the read handle falls back to a + * key-based lookup at the live file.
    • + *
  • {@code DELETED}: excluded entirely.
    • + *
+ * + *

Builds no SQL strings and no temporary DataFrames. {@code baseFilePath} is left null here; the + * read handle resolves the concrete base file from the snapshot's file slice for the fileId. + */ +public final class DefaultVectorFetchPlanner implements VectorFetchPlanner { + + private static final long serialVersionUID = 1L; + private static final char KEY_SEP = '\u0001'; + + @Override + public HoodieData plan(HoodieData candidates, + VectorSearchSnapshot snapshot, + HoodieEngineContext engineContext) { + return candidates + .flatMapToPair(DefaultVectorFetchPlanner::toFileKeyed) + .groupByKey() + .map(entry -> buildTask(entry.getKey(), entry.getValue())); + } + + /** DELETED candidates and those without a live location are dropped (empty iterator). */ + private static Iterator> toFileKeyed(ArbitratedVectorCandidate c) { + if (!c.isServable() || c.getLiveLocation() == null) { + return Collections.emptyIterator(); + } + HoodieRecordGlobalLocation loc = c.getLiveLocation(); + String key = loc.getPartitionPath() + KEY_SEP + loc.getFileId(); + return Collections.singletonList(Pair.of(key, c)).iterator(); + } + + private static VectorFetchTask buildTask(String fileKey, Iterable group) { + List requests = new ArrayList<>(); + String partitionPath = null; + String fileId = null; + String baseInstant = null; + for (ArbitratedVectorCandidate c : group) { + HoodieRecordGlobalLocation loc = c.getLiveLocation(); + if (partitionPath == null) { + partitionPath = loc.getPartitionPath(); + fileId = loc.getFileId(); + baseInstant = loc.getInstantTime(); + } + VectorCandidate cand = c.getCandidate(); + boolean serve = c.getState() == VectorCandidateState.SERVE; + long rowPosition = serve && cand.getPostingLocator() != null + ? cand.getPostingLocator().getRowPosition() : -1L; + requests.add(new VectorRowRequest( + cand.getRecordKey(), rowPosition, c.getState(), cand.getApproximateDistance())); + } + return new VectorFetchTask(partitionPath, fileId, null, baseInstant, requests); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ExactVectorScorer.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ExactVectorScorer.java new file mode 100644 index 0000000000000..b9df4facff0f8 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ExactVectorScorer.java @@ -0,0 +1,33 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.index.vector.VectorDistanceMetric; + +import java.io.Serializable; + +/** + * Computes the exact metric distance between a query and a full-precision candidate vector + * (RFC-109 §10). Accumulates in float64 and keeps squared L2 internally; the surfaced value + * follows the requested {@link VectorDistanceMetric}. + */ +public interface ExactVectorScorer extends Serializable { + + double distance(float[] query, float[] candidate, VectorDistanceMetric metric); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/HoodieVectorBatchReadHandle.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/HoodieVectorBatchReadHandle.java new file mode 100644 index 0000000000000..d1ebaff1df180 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/HoodieVectorBatchReadHandle.java @@ -0,0 +1,45 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.util.Iterator; + +/** + * A Hudi read handle for projected, position-based base-file reads (RFC-109 §9). Given a + * {@link VectorFetchTask}, returns only the record-key and vector columns as {@link VectorRecord}s + * — no full-row materialization. SERVE rows are read by row position (page-index skipping); STALE + * rows fall back to a key-based lookup within the same file. + * + *

The interface is engine- and format-neutral: implementations group positions by row group and + * page, cache footer/page-index metadata, coalesce neighboring ranges, use bounded concurrency, and + * never expose storage-specific (GCS/S3) APIs to the vector-search layer. The Parquet implementation + * is the first; the same shape supports ORC and is reusable for record-index point reads. + */ +public interface HoodieVectorBatchReadHandle extends AutoCloseable { + + /** + * Read the requested rows of one file, decoding only {@code recordKeyField} and {@code vectorColumn}. + * + * @param task the per-file fetch task (base file + row requests) + * @param recordKeyField the record-key column name + * @param vectorColumn the vector column name + * @return an iterator of decoded {@link VectorRecord}s (record key + vector + live location) + */ + Iterator read(VectorFetchTask task, String recordKeyField, String vectorColumn); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/HoodieVectorBatchReadHandleSupplier.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/HoodieVectorBatchReadHandleSupplier.java new file mode 100644 index 0000000000000..fe97d56fdebe0 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/HoodieVectorBatchReadHandleSupplier.java @@ -0,0 +1,33 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; + +/** + * Serializable factory for a {@link HoodieVectorBatchReadHandle} (RFC-109 §9, §10). The reranker + * creates one handle per partition/task-runner via this supplier, so the concrete (Parquet) handle + * — which lives in an engine/format module — is constructed on the executor without the common + * reranker depending on it. Engine adapters provide the implementation. + */ +@FunctionalInterface +public interface HoodieVectorBatchReadHandleSupplier extends Serializable { + + HoodieVectorBatchReadHandle get(); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ListVectorCandidatePool.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ListVectorCandidatePool.java new file mode 100644 index 0000000000000..f4e18f27c41aa --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ListVectorCandidatePool.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.data.HoodieListData; + +import java.util.List; + +/** In-memory retained pool for local execution and engine-neutral tests. */ +public final class ListVectorCandidatePool implements VectorCandidatePool { + + private static final long serialVersionUID = 1L; + + private final VectorContinuationController controller; + + public ListVectorCandidatePool(List orderedCandidates, VectorSearchBudget budget) { + this.controller = new VectorContinuationController<>(orderedCandidates, + budget.getInitialRerankCandidates(), budget.getRerankBatchSize(), budget.getMaxRerankCandidates()); + } + + @Override + public boolean hasMore() { + return controller.hasMore(); + } + + @Override + public HoodieData nextBatch() { + return HoodieListData.eager(controller.nextBatch()); + } + + @Override + public int consumed() { + return controller.consumed(); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/RecordIndexLookup.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/RecordIndexLookup.java new file mode 100644 index 0000000000000..08c4d7fef3d34 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/RecordIndexLookup.java @@ -0,0 +1,40 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import java.io.Serializable; +import java.util.List; +import java.util.Map; + +/** + * A snapshot-pinned batched Record-Level Index lookup (RFC-109 §7). Given a batch of record keys, + * returns the current live location for each key that exists at the pinned table instant; keys with + * no entry (absent from the map) are treated as deleted. + * + *

Engine adapters supply the concrete implementation (wrapping + * {@code readRecordIndexLocationsWithKeys} at {@code snapshot.tableInstant}). Keeping it an injected + * SAM lets the arbiter logic stay engine-neutral and unit-testable with a fake lookup. + */ +@FunctionalInterface +public interface RecordIndexLookup extends Serializable { + + Map lookup(List recordKeys, String tableInstant); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/RecordIndexVectorCandidateArbiter.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/RecordIndexVectorCandidateArbiter.java new file mode 100644 index 0000000000000..e5821e85d7d06 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/RecordIndexVectorCandidateArbiter.java @@ -0,0 +1,159 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.engine.HoodieEngineContext; +import org.apache.hudi.common.index.vector.VectorIndexArbiter; +import org.apache.hudi.common.index.vector.VectorStalePolicy; +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; +import org.apache.hudi.exception.HoodieException; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.ArrayList; +import java.util.Iterator; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Snapshot-aware RLI candidate arbiter (RFC-109 §7). Wraps the pure decision core + * {@link VectorIndexArbiter#classify} with a batched, snapshot-pinned {@link RecordIndexLookup}: + * per partition it collects candidate keys, performs a single batched RLI lookup, and classifies + * each candidate against its current live location. + * + *

    + *
  • {@code SERVE}: current location matches the posting locator; positional trust preserved.
    • + *
  • {@code STALE}: current location differs; preserved for exact-mode key fallback.
    • + *
  • {@code DELETED}: RLI miss; dropped from the output.
    • + *
+ * + *

Posting locations are treated as hints until this arbitration. The lookup must be pinned to + * {@code snapshot.tableInstant} by the caller so MDT, RLI, file-slice, and base reads share one + * instant. + */ +public final class RecordIndexVectorCandidateArbiter implements VectorCandidateArbiter { + + private static final long serialVersionUID = 1L; + private static final Logger LOG = LoggerFactory.getLogger(RecordIndexVectorCandidateArbiter.class); + private static final int DEFAULT_LOOKUP_BATCH_SIZE = 1024; + + private final RecordIndexLookup lookup; + private final int lookupBatchSize; + + public RecordIndexVectorCandidateArbiter(RecordIndexLookup lookup) { + this(lookup, DEFAULT_LOOKUP_BATCH_SIZE); + } + + public RecordIndexVectorCandidateArbiter(RecordIndexLookup lookup, int lookupBatchSize) { + if (lookupBatchSize <= 0) { + throw new IllegalArgumentException("lookupBatchSize must be positive"); + } + this.lookup = lookup; + this.lookupBatchSize = lookupBatchSize; + } + + @Override + public HoodieData arbitrate(HoodieData candidates, + VectorSearchRequest request, + VectorSearchSnapshot snapshot, + HoodieEngineContext engineContext) { + RecordIndexLookup rli = this.lookup; + int batchSize = this.lookupBatchSize; + return candidates.mapPartitions( + it -> arbitratePartition(it, rli, snapshot.getTableInstant(), request.getStalePolicy(), batchSize), true); + } + + private static Iterator arbitratePartition( + Iterator candidates, + RecordIndexLookup rli, + String tableInstant, + VectorStalePolicy stalePolicy, + int batchSize) { + List out = new ArrayList<>(); + List batch = new ArrayList<>(batchSize); + int staleCount = 0; + while (candidates.hasNext()) { + batch.add(candidates.next()); + if (batch.size() == batchSize) { + staleCount += arbitrateBatch(batch, rli, tableInstant, stalePolicy, out); + batch.clear(); + } + } + if (!batch.isEmpty()) { + staleCount += arbitrateBatch(batch, rli, tableInstant, stalePolicy, out); + } + if (staleCount > 0 && stalePolicy == VectorStalePolicy.WARN) { + LOG.warn("Vector index contained {} stale candidates at table instant {}; using RLI fallback locations", + staleCount, tableInstant); + } + return out.iterator(); + } + + private static int arbitrateBatch( + List batch, + RecordIndexLookup rli, + String tableInstant, + VectorStalePolicy stalePolicy, + List out) { + Set keys = new LinkedHashSet<>(); + for (VectorCandidate candidate : batch) { + keys.add(candidate.getRecordKey()); + } + Map current = + rli.lookup(new ArrayList<>(keys), tableInstant); + int staleCount = 0; + for (VectorCandidate candidate : batch) { + HoodieRecordGlobalLocation live = current.get(candidate.getRecordKey()); + VectorPostingLocator locator = candidate.getPostingLocator(); + VectorCandidateState state = toState(VectorIndexArbiter.classify( + locator == null ? null : locator.getPartitionPath(), + locator == null ? null : locator.getFileId(), + locator == null ? null : locator.getBaseInstant(), + live)); + if (state == VectorCandidateState.DELETED) { + continue; + } + if (state == VectorCandidateState.STALE) { + staleCount++; + if (stalePolicy == VectorStalePolicy.FAIL) { + throw new HoodieException("Stale vector candidate '" + candidate.getRecordKey() + + "' at table instant " + tableInstant); + } + } + out.add(new ArbitratedVectorCandidate(candidate, state, live)); + } + return staleCount; + } + + private static VectorCandidateState toState(VectorIndexArbiter.Decision decision) { + switch (decision) { + case SERVE: + return VectorCandidateState.SERVE; + case STALE: + return VectorCandidateState.STALE; + case DELETED: + default: + return VectorCandidateState.DELETED; + } + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ThresholdVectorExecutionModeSelector.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ThresholdVectorExecutionModeSelector.java new file mode 100644 index 0000000000000..4c3815fb06f53 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/ThresholdVectorExecutionModeSelector.java @@ -0,0 +1,52 @@ +/* + * 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.hudi.common.index.vector.search; + +/** + * Normative candidate-threshold execution selector (RFC-109 §11A, selector version + * {@code candidate-threshold-v1}). Pure function of the request's requested mode and budget; + * no engine or runtime state, so LOCAL/DISTRIBUTED selection is deterministic and reproducible. + */ +public final class ThresholdVectorExecutionModeSelector implements VectorExecutionModeSelector { + + private static final long serialVersionUID = 1L; + + @Override + public VectorExecutionDecision select(VectorSearchRequest request) { + VectorSearchBudget budget = request.getBudget(); + VectorExecutionMode requested = budget.getRequestedExecutionMode(); + int maxRerank = budget.getMaxRerankCandidates(); + int threshold = budget.getLocalExecutionThreshold(); + + VectorExecutionMode selected; + switch (requested) { + case LOCAL: + selected = VectorExecutionMode.LOCAL; + break; + case DISTRIBUTED: + selected = VectorExecutionMode.DISTRIBUTED; + break; + case AUTO: + default: + selected = maxRerank <= threshold ? VectorExecutionMode.LOCAL : VectorExecutionMode.DISTRIBUTED; + break; + } + return new VectorExecutionDecision(requested, selected, maxRerank, threshold, SELECTOR_VERSION); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidate.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidate.java new file mode 100644 index 0000000000000..e3063d8d10b4d --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidate.java @@ -0,0 +1,71 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; + +/** + * A retained ANN candidate emitted by a {@link VectorCandidateSource} (RFC-109 §4): the logical + * record key, its cluster/shard, the approximate (squared L2) distance from RaBitQ scoring, and the + * posting locator hint. Record keys and locators are decoded only for retained candidates, never + * for rejected posting rows. + */ +public final class VectorCandidate implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String recordKey; + private final int clusterId; + private final int shardId; + private final double approximateDistance; + private final VectorPostingLocator postingLocator; + + public VectorCandidate(String recordKey, + int clusterId, + int shardId, + double approximateDistance, + VectorPostingLocator postingLocator) { + this.recordKey = recordKey; + this.clusterId = clusterId; + this.shardId = shardId; + this.approximateDistance = approximateDistance; + this.postingLocator = postingLocator; + } + + public String getRecordKey() { + return recordKey; + } + + public int getClusterId() { + return clusterId; + } + + public int getShardId() { + return shardId; + } + + /** Approximate squared-L2 distance from RaBitQ scoring; kept squared internally through ranking. */ + public double getApproximateDistance() { + return approximateDistance; + } + + public VectorPostingLocator getPostingLocator() { + return postingLocator; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateArbiter.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateArbiter.java new file mode 100644 index 0000000000000..6e9b01efccea0 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateArbiter.java @@ -0,0 +1,39 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.engine.HoodieEngineContext; + +import java.io.Serializable; + +/** + * Validates candidate freshness against the Record-Level Index at the pinned snapshot (RFC-109 §7). + * Treats posting locations as hints until arbitration, preserves STALE candidates for exact-mode key + * fallback, drops DELETED candidates, and uses the same table instant as the MDT/file-slice/base reads. + * The pure decision core stays in {@code VectorIndexArbiter.classify}; implementations add the + * snapshot-aware distributed RLI lookup. + */ +public interface VectorCandidateArbiter extends Serializable { + + HoodieData arbitrate(HoodieData candidates, + VectorSearchRequest request, + VectorSearchSnapshot snapshot, + HoodieEngineContext engineContext); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateOverlay.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateOverlay.java new file mode 100644 index 0000000000000..997b798a734e9 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateOverlay.java @@ -0,0 +1,68 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector.search; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** Resolves packed-block candidates with the canonical delta overlay before finalist arbitration. */ +public final class VectorCandidateOverlay { + + private static final Comparator ORDER = Comparator + .comparingDouble(VectorCandidate::getApproximateDistance) + .thenComparing(VectorCandidate::getRecordKey); + + private VectorCandidateOverlay() { + } + + /** + * Retains base candidates with overlay slack, then applies canonical delta precedence and + * tombstone suppression. A live delta always replaces the packed-block row for the same key, + * even when its approximate distance is worse; otherwise an update could resurrect stale code. + */ + public static List resolve( + Collection baseCandidates, + Collection deltaCandidates, + Collection tombstonedDeltaKeys, + int maxCandidates, + int overlaySlack) { + if (maxCandidates < 0 || overlaySlack < 0) { + throw new IllegalArgumentException("candidate bounds must be non-negative"); + } + int retainedBaseCount = Math.min(baseCandidates.size(), Math.addExact(maxCandidates, overlaySlack)); + List orderedBase = new ArrayList<>(baseCandidates); + orderedBase.sort(ORDER); + + Set tombstones = new HashSet<>(tombstonedDeltaKeys); + Map resolved = new HashMap<>(); + for (int i = 0; i < retainedBaseCount; i++) { + VectorCandidate candidate = orderedBase.get(i); + if (!tombstones.contains(candidate.getRecordKey())) { + resolved.put(candidate.getRecordKey(), candidate); + } + } + for (VectorCandidate delta : deltaCandidates) { + if (tombstones.contains(delta.getRecordKey())) { + resolved.remove(delta.getRecordKey()); + } else { + resolved.put(delta.getRecordKey(), delta); + } + } + + List result = new ArrayList<>(resolved.values()); + result.sort(ORDER); + if (result.size() > maxCandidates) { + return new ArrayList<>(result.subList(0, maxCandidates)); + } + return result; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidatePool.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidatePool.java new file mode 100644 index 0000000000000..eb9bf067c9d25 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidatePool.java @@ -0,0 +1,20 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; + +import java.io.Serializable; + +/** A single-scan, retained candidate pool consumed in ordered continuation windows. */ +public interface VectorCandidatePool extends Serializable { + + boolean hasMore(); + + HoodieData nextBatch(); + + int consumed(); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateSource.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateSource.java new file mode 100644 index 0000000000000..3bfcb31583146 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateSource.java @@ -0,0 +1,35 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.engine.HoodieEngineContext; + +import java.io.Serializable; + +/** + * Produces ANN candidates for a plan (RFC-109 §4). The MDT implementation owns posting decoding, + * overlay resolution, pass-1 filtering, pass-2 scoring, and bounded candidate retention — and it + * MUST NOT invoke Spark SQL, file-format readers, or exact-read code. It decodes record keys and + * locators only for retained candidates and returns at most {@code maxRerankCandidates} ordered by + * approximate distance, in a single scan. + */ +public interface VectorCandidateSource extends Serializable { + + VectorCandidatePool scan(VectorSearchPlan plan, HoodieEngineContext engineContext); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateState.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateState.java new file mode 100644 index 0000000000000..1b942f5952e7a --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorCandidateState.java @@ -0,0 +1,36 @@ +/* + * 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.hudi.common.index.vector.search; + +/** + * Freshness verdict for a finalist candidate, produced by RLI arbitration (RFC-109 §7). + * Engine-neutral successor to the internal arbiter decision enum. + * + *

    + *
  • {@code SERVE}: posting still faithfully represents the live record; positional trust holds.
    • + *
  • {@code STALE}: record was rewritten/moved; posting locator is a stale hint, exact mode must + * key-fetch at the live location.
    • + *
  • {@code DELETED}: record no longer exists; drop.
    • + *
+ */ +public enum VectorCandidateState { + SERVE, + STALE, + DELETED +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorContinuationController.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorContinuationController.java new file mode 100644 index 0000000000000..554a57980caaa --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorContinuationController.java @@ -0,0 +1,81 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.util.List; + +/** + * Windows a single retained, distance-ordered candidate pool into continuation batches + * (RFC-109 §10). The candidate scan retains one ordered pool of at most + * {@code maxRerankCandidates} in a single MDT scan; this controller draws successive batches + * from that retained pool — it never rescans MDT postings. + * + *

The first {@link #nextBatch()} returns up to {@code initialRerankCandidates}; subsequent calls + * return up to {@code rerankBatchSize}, until the pool (bounded by {@code maxRerankCandidates}) is + * exhausted. The reranker calls {@link #nextBatch()} while {@link #hasMore()} and the top-K + * accumulator still {@link VectorTopKAccumulator#needsMore() needs more} live results (and the + * deadline has not passed). + */ +public final class VectorContinuationController { + + private final List orderedPool; + private final int initialRerankCandidates; + private final int rerankBatchSize; + private final int effectiveMax; + private int cursor; + + public VectorContinuationController(List orderedPool, + int initialRerankCandidates, + int rerankBatchSize, + int maxRerankCandidates) { + if (initialRerankCandidates <= 0 || rerankBatchSize <= 0) { + throw new IllegalArgumentException("batch sizes must be positive"); + } + this.orderedPool = orderedPool; + this.initialRerankCandidates = initialRerankCandidates; + this.rerankBatchSize = rerankBatchSize; + this.effectiveMax = Math.min(orderedPool.size(), Math.max(0, maxRerankCandidates)); + this.cursor = 0; + } + + /** Whether more retained candidates remain to draw (within {@code maxRerankCandidates}). */ + public boolean hasMore() { + return cursor < effectiveMax; + } + + /** Number of candidates drawn so far (monotonic; never exceeds the retained pool bound). */ + public int consumed() { + return cursor; + } + + /** + * Draw the next continuation batch as a window over the retained pool. First call returns up to + * {@code initialRerankCandidates}; later calls up to {@code rerankBatchSize}. Never rescans. + */ + public List nextBatch() { + if (!hasMore()) { + return java.util.Collections.emptyList(); + } + int size = cursor == 0 ? initialRerankCandidates : rerankBatchSize; + int end = Math.min(cursor + size, effectiveMax); + List batch = orderedPool.subList(cursor, end); + cursor = end; + return batch; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExactReranker.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExactReranker.java new file mode 100644 index 0000000000000..5456075f84ed1 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExactReranker.java @@ -0,0 +1,43 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.engine.HoodieEngineContext; + +import java.io.Serializable; + +/** + * Reads candidate vectors through Hudi read handles and produces exact top-K results (RFC-109 §10). + * + *

Continuation: the candidate scan already retained one ordered pool of at most + * {@code maxRerankCandidates}. The reranker starts with the initial rerank batch, and continues to + * draw successive batches from that retained pool when stale/deleted records leave fewer + * than K live results — it MUST NOT rescan MDT postings per batch. It stops when K live exact + * results exist, the retained pool is exhausted, or the deadline is reached, returning an explicit + * {@link VectorSearchStatus} rather than silently returning fewer than K. Executes LOCAL or + * DISTRIBUTED per the plan's {@link VectorExecutionDecision}. + */ +public interface VectorExactReranker extends Serializable { + + HoodieData rerank(HoodieData tasks, + VectorSearchRequest request, + VectorSearchSnapshot snapshot, + HoodieEngineContext engineContext); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExecutionDecision.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExecutionDecision.java new file mode 100644 index 0000000000000..dd1b1ba1fc0e9 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExecutionDecision.java @@ -0,0 +1,78 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; + +/** + * The recorded outcome of execution-locality selection (RFC-109 §11A). Carries both the + * requested and selected mode plus the inputs to the decision so it can be emitted verbatim into + * query metrics and workload-profile results. + */ +public final class VectorExecutionDecision implements Serializable { + + private static final long serialVersionUID = 1L; + + private final VectorExecutionMode requestedMode; + private final VectorExecutionMode selectedMode; + private final int maxRerankCandidates; + private final int localExecutionThreshold; + private final String selectorVersion; + + public VectorExecutionDecision(VectorExecutionMode requestedMode, + VectorExecutionMode selectedMode, + int maxRerankCandidates, + int localExecutionThreshold, + String selectorVersion) { + this.requestedMode = requestedMode; + this.selectedMode = selectedMode; + this.maxRerankCandidates = maxRerankCandidates; + this.localExecutionThreshold = localExecutionThreshold; + this.selectorVersion = selectorVersion; + } + + public VectorExecutionMode getRequestedMode() { + return requestedMode; + } + + public VectorExecutionMode getSelectedMode() { + return selectedMode; + } + + public int getMaxRerankCandidates() { + return maxRerankCandidates; + } + + public int getLocalExecutionThreshold() { + return localExecutionThreshold; + } + + public String getSelectorVersion() { + return selectorVersion; + } + + @Override + public String toString() { + return "VectorExecutionDecision{requested=" + requestedMode + + ", selected=" + selectedMode + + ", maxRerankCandidates=" + maxRerankCandidates + + ", localThreshold=" + localExecutionThreshold + + ", selectorVersion=" + selectorVersion + '}'; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExecutionMode.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExecutionMode.java new file mode 100644 index 0000000000000..6ef77441e6c82 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExecutionMode.java @@ -0,0 +1,31 @@ +/* + * 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.hudi.common.index.vector.search; + +/** + * Requested/selected execution locality for a vector search. Engine-neutral: {@code LOCAL} means + * the common executor runs the bounded fetch/score tasks through a local task runner (e.g. over + * {@code HoodieListData}); {@code DISTRIBUTED} means it schedules them on the engine's cluster. + * {@code AUTO} defers to the normative selection rule (RFC-109 §11A). + */ +public enum VectorExecutionMode { + AUTO, + LOCAL, + DISTRIBUTED +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExecutionModeSelector.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExecutionModeSelector.java new file mode 100644 index 0000000000000..268b1d499625b --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorExecutionModeSelector.java @@ -0,0 +1,38 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; + +/** + * Chooses LOCAL vs DISTRIBUTED execution for a vector search (RFC-109 §11A). + * + *

The default rule is normative and fixed for selector version {@code candidate-threshold-v1}: + * an explicit {@code LOCAL}/{@code DISTRIBUTED} request is honored verbatim, and {@code AUTO} + * selects {@code LOCAL} iff {@code maxRerankCandidates <= localExecutionThreshold} (default 8192). + * No candidate-count, file-count, byte-estimate, or engine heuristic may silently alter the + * {@code AUTO} rule in this version; a future adaptive selector requires a new selector version. + */ +public interface VectorExecutionModeSelector extends Serializable { + + int DEFAULT_LOCAL_EXECUTION_THRESHOLD = VectorSearchBudget.DEFAULT_LOCAL_EXECUTION_THRESHOLD; + String SELECTOR_VERSION = "candidate-threshold-v1"; + + VectorExecutionDecision select(VectorSearchRequest request); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorFetchPlanner.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorFetchPlanner.java new file mode 100644 index 0000000000000..adeb6edef78af --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorFetchPlanner.java @@ -0,0 +1,36 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.engine.HoodieEngineContext; + +import java.io.Serializable; + +/** + * Groups arbitrated candidates into per-file-slice fetch tasks (RFC-109 §8). Resolves file slices + * against the pinned snapshot, preserves row positions for {@code SERVE}, plans key-based fallback + * for {@code STALE}, and never builds SQL strings or temporary DataFrames. + */ +public interface VectorFetchPlanner extends Serializable { + + HoodieData plan(HoodieData candidates, + VectorSearchSnapshot snapshot, + HoodieEngineContext engineContext); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorFetchTask.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorFetchTask.java new file mode 100644 index 0000000000000..129da1b5d2ab7 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorFetchTask.java @@ -0,0 +1,76 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; +import java.util.Collections; +import java.util.List; + +/** + * A batch of rows to read from a single snapshot-resolved base file slice (RFC-109 §8). Produced + * by the {@link VectorFetchPlanner} by grouping arbitrated candidates by file slice, so the read + * handle can coalesce positions within one file/row-group/page. Compact and serializable — carries + * only paths and row requests, never engine or SQL types. + */ +public final class VectorFetchTask implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String partitionPath; + private final String fileId; + private final String baseFilePath; + private final String baseInstant; + private final List requests; + + public VectorFetchTask(String partitionPath, + String fileId, + String baseFilePath, + String baseInstant, + List requests) { + this.partitionPath = partitionPath; + this.fileId = fileId; + this.baseFilePath = baseFilePath; + this.baseInstant = baseInstant; + this.requests = requests == null ? Collections.emptyList() : requests; + } + + public String getPartitionPath() { + return partitionPath; + } + + public String getFileId() { + return fileId; + } + + public String getBaseFilePath() { + return baseFilePath; + } + + public String getBaseInstant() { + return baseInstant; + } + + public List getRequests() { + return requests; + } + + public int size() { + return requests.size(); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorIndexSnapshot.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorIndexSnapshot.java new file mode 100644 index 0000000000000..d45c59634c753 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorIndexSnapshot.java @@ -0,0 +1,70 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; + +/** + * Immutable identity of the active vector-index generation used to serve a query (RFC-109). + * Every field is versioned through the manifest so readers can reject unsupported or mismatched + * encodings rather than silently mis-scoring. Pinned for the whole request alongside the table + * instant in {@link VectorSearchSnapshot}. + */ +public final class VectorIndexSnapshot implements Serializable { + + private static final long serialVersionUID = 1L; + + private final int generationId; + private final int factorVersion; + private final int blockFormatVersion; + private final String rotationVersion; + private final String quantizerVersion; + + public VectorIndexSnapshot(int generationId, + int factorVersion, + int blockFormatVersion, + String rotationVersion, + String quantizerVersion) { + this.generationId = generationId; + this.factorVersion = factorVersion; + this.blockFormatVersion = blockFormatVersion; + this.rotationVersion = rotationVersion; + this.quantizerVersion = quantizerVersion; + } + + public int getGenerationId() { + return generationId; + } + + public int getFactorVersion() { + return factorVersion; + } + + public int getBlockFormatVersion() { + return blockFormatVersion; + } + + public String getRotationVersion() { + return rotationVersion; + } + + public String getQuantizerVersion() { + return quantizerVersion; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorPostingLocator.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorPostingLocator.java new file mode 100644 index 0000000000000..6af590d61655b --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorPostingLocator.java @@ -0,0 +1,110 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; + +/** + * Physical hint for where a candidate's packed posting row lives (RFC-109 §4). It combines the + * logical index coordinates (generation/cluster/shard/block/ordinal) with an optional data-table + * location hint (partition/file/rowPosition) decoded from the posting. + * + *

Per the architectural rule, the data-table location is a hint only: it is trusted for + * a positional read solely after RLI arbitration returns {@link VectorCandidateState#SERVE}. When + * the base instant recorded here diverges from the live RLI location the candidate is STALE and the + * fetch planner falls back to a key-based read. + */ +public final class VectorPostingLocator implements Serializable { + + private static final long serialVersionUID = 1L; + + // Logical index coordinates. + private final int generationId; + private final int clusterId; + private final int shardId; + private final long blockId; + private final int vectorOrdinal; + + // Data-table location hint (may be null when only key-fallback is possible). + private final String partitionPath; + private final String fileId; + private final String baseInstant; + private final long rowPosition; + + public VectorPostingLocator(int generationId, + int clusterId, + int shardId, + long blockId, + int vectorOrdinal, + String partitionPath, + String fileId, + String baseInstant, + long rowPosition) { + this.generationId = generationId; + this.clusterId = clusterId; + this.shardId = shardId; + this.blockId = blockId; + this.vectorOrdinal = vectorOrdinal; + this.partitionPath = partitionPath; + this.fileId = fileId; + this.baseInstant = baseInstant; + this.rowPosition = rowPosition; + } + + public int getGenerationId() { + return generationId; + } + + public int getClusterId() { + return clusterId; + } + + public int getShardId() { + return shardId; + } + + public long getBlockId() { + return blockId; + } + + public int getVectorOrdinal() { + return vectorOrdinal; + } + + public String getPartitionPath() { + return partitionPath; + } + + public String getFileId() { + return fileId; + } + + public String getBaseInstant() { + return baseInstant; + } + + public long getRowPosition() { + return rowPosition; + } + + /** Whether a positional data-table hint is present (still subject to RLI arbitration). */ + public boolean hasDataLocationHint() { + return fileId != null && rowPosition >= 0; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorRecord.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorRecord.java new file mode 100644 index 0000000000000..1ad49389b36de --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorRecord.java @@ -0,0 +1,56 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import java.io.Serializable; + +/** + * A decoded record read back from the base table by a {@link org.apache.hudi.common.index.vector.search} + * read handle (RFC-109 §9): the logical record key, its full-precision vector, and the live + * location it was read from. Only the record-key and vector columns are decoded — no full-row + * materialization. + */ +public final class VectorRecord implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String recordKey; + private final float[] vector; + private final HoodieRecordGlobalLocation location; + + public VectorRecord(String recordKey, float[] vector, HoodieRecordGlobalLocation location) { + this.recordKey = recordKey; + this.vector = vector; + this.location = location; + } + + public String getRecordKey() { + return recordKey; + } + + public float[] getVector() { + return vector; + } + + public HoodieRecordGlobalLocation getLocation() { + return location; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorRowRequest.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorRowRequest.java new file mode 100644 index 0000000000000..6a089780055c5 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorRowRequest.java @@ -0,0 +1,66 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; + +/** + * One row to fetch within a {@link VectorFetchTask} (RFC-109 §8). For {@link VectorCandidateState#SERVE} + * the {@link #rowPosition} is authoritative for a positional read; for {@link VectorCandidateState#STALE} + * the position is ignored and the read handle falls back to a key-based lookup. + */ +public final class VectorRowRequest implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String recordKey; + private final long rowPosition; + private final VectorCandidateState state; + private final double approximateDistance; + + public VectorRowRequest(String recordKey, + long rowPosition, + VectorCandidateState state, + double approximateDistance) { + this.recordKey = recordKey; + this.rowPosition = rowPosition; + this.state = state; + this.approximateDistance = approximateDistance; + } + + public String getRecordKey() { + return recordKey; + } + + public long getRowPosition() { + return rowPosition; + } + + public VectorCandidateState getState() { + return state; + } + + public double getApproximateDistance() { + return approximateDistance; + } + + public boolean isPositional() { + return state == VectorCandidateState.SERVE && rowPosition >= 0; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchBudget.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchBudget.java new file mode 100644 index 0000000000000..b8e1031b3fdc9 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchBudget.java @@ -0,0 +1,127 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; +import java.util.Objects; + +/** + * Per-request resource and continuation budget for a vector search (RFC-109 §1). + * + *

Continuation semantics: the candidate scan retains one ordered pool of at most + * {@link #maxRerankCandidates} in a single MDT scan; exact rerank consumes it in batches of + * {@link #rerankBatchSize} starting at {@link #initialRerankCandidates}, drawing more only when + * stale/deleted finalists leave fewer than K live results. Continuation is a window over the + * retained pool, never a re-scan. + * + *

Execution locality: {@link #requestedExecutionMode} plus {@link #localExecutionThreshold} + * feed the normative selector (§11A) — {@code AUTO} selects {@code LOCAL} iff + * {@code maxRerankCandidates <= localExecutionThreshold}. + */ +public final class VectorSearchBudget implements Serializable { + + private static final long serialVersionUID = 1L; + + /** Default local-vs-distributed candidate threshold (RFC-109 §11A). */ + public static final int DEFAULT_LOCAL_EXECUTION_THRESHOLD = 8192; + public static final int DEFAULT_INITIAL_RERANK_CANDIDATES = 256; + public static final int DEFAULT_RERANK_BATCH_SIZE = 128; + public static final int DEFAULT_MAX_RERANK_CANDIDATES = 4096; + + private final long timeoutMs; + private final int initialRerankCandidates; + private final int rerankBatchSize; + private final int maxRerankCandidates; + private final int maxFetchTasks; + private final int maxFetchConcurrency; + private final VectorExecutionMode requestedExecutionMode; + private final int localExecutionThreshold; + private final DeadlinePolicy deadlinePolicy; + + public VectorSearchBudget(long timeoutMs, + int initialRerankCandidates, + int rerankBatchSize, + int maxRerankCandidates, + int maxFetchTasks, + int maxFetchConcurrency, + VectorExecutionMode requestedExecutionMode, + int localExecutionThreshold, + DeadlinePolicy deadlinePolicy) { + this.timeoutMs = timeoutMs; + this.initialRerankCandidates = initialRerankCandidates; + this.rerankBatchSize = rerankBatchSize; + this.maxRerankCandidates = maxRerankCandidates; + this.maxFetchTasks = maxFetchTasks; + this.maxFetchConcurrency = maxFetchConcurrency; + this.requestedExecutionMode = Objects.requireNonNull(requestedExecutionMode, "requestedExecutionMode"); + this.localExecutionThreshold = localExecutionThreshold; + this.deadlinePolicy = Objects.requireNonNull(deadlinePolicy, "deadlinePolicy"); + } + + /** A sensible default budget for {@code topK}: AUTO execution, partial-on-deadline off (FAIL). */ + public static VectorSearchBudget defaults(int topK, long timeoutMs) { + int maxRerank = Math.max(DEFAULT_MAX_RERANK_CANDIDATES, topK * 16); + return new VectorSearchBudget( + timeoutMs, + Math.max(DEFAULT_INITIAL_RERANK_CANDIDATES, topK), + DEFAULT_RERANK_BATCH_SIZE, + maxRerank, + Integer.MAX_VALUE, + Math.max(1, Runtime.getRuntime().availableProcessors()), + VectorExecutionMode.AUTO, + DEFAULT_LOCAL_EXECUTION_THRESHOLD, + DeadlinePolicy.FAIL); + } + + public long getTimeoutMs() { + return timeoutMs; + } + + public int getInitialRerankCandidates() { + return initialRerankCandidates; + } + + public int getRerankBatchSize() { + return rerankBatchSize; + } + + public int getMaxRerankCandidates() { + return maxRerankCandidates; + } + + public int getMaxFetchTasks() { + return maxFetchTasks; + } + + public int getMaxFetchConcurrency() { + return maxFetchConcurrency; + } + + public VectorExecutionMode getRequestedExecutionMode() { + return requestedExecutionMode; + } + + public int getLocalExecutionThreshold() { + return localExecutionThreshold; + } + + public DeadlinePolicy getDeadlinePolicy() { + return deadlinePolicy; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchExecutor.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchExecutor.java new file mode 100644 index 0000000000000..9cd6dc714c429 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchExecutor.java @@ -0,0 +1,36 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.engine.HoodieEngineContext; + +import java.io.Serializable; + +/** + * The single engine-neutral entry point for vector search (RFC-109 §11). Pins one snapshot, + * probes IVF clusters, scans MDT postings, reduces the candidate pool, RLI-arbitrates, plans + * file-slice fetches, chooses LOCAL/DISTRIBUTED execution, performs projected positional/key reads, + * scores exactly, and reduces to top-K — all under one request deadline. Never invokes + * {@code spark.sql(...)} and never reconstructs SQL/DataFrames to execute exact fetches. + */ +public interface VectorSearchExecutor extends Serializable { + + HoodieData execute(VectorSearchRequest request, HoodieEngineContext engineContext); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchPlan.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchPlan.java new file mode 100644 index 0000000000000..5ff275cc580c8 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchPlan.java @@ -0,0 +1,57 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; +import java.util.Objects; + +/** + * The resolved, engine-neutral plan for a single vector search: the immutable request, the pinned + * {@link VectorSearchSnapshot}, and the {@link VectorExecutionDecision} chosen by the selector. + * Built once by the orchestrator and threaded through every stage so all stages share one snapshot + * and one execution decision (RFC-109 §11). + */ +public final class VectorSearchPlan implements Serializable { + + private static final long serialVersionUID = 1L; + + private final VectorSearchRequest request; + private final VectorSearchSnapshot snapshot; + private final VectorExecutionDecision executionDecision; + + public VectorSearchPlan(VectorSearchRequest request, + VectorSearchSnapshot snapshot, + VectorExecutionDecision executionDecision) { + this.request = Objects.requireNonNull(request, "request"); + this.snapshot = Objects.requireNonNull(snapshot, "snapshot"); + this.executionDecision = Objects.requireNonNull(executionDecision, "executionDecision"); + } + + public VectorSearchRequest getRequest() { + return request; + } + + public VectorSearchSnapshot getSnapshot() { + return snapshot; + } + + public VectorExecutionDecision getExecutionDecision() { + return executionDecision; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchRequest.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchRequest.java new file mode 100644 index 0000000000000..1943788a6753c --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchRequest.java @@ -0,0 +1,124 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.index.vector.VectorDistanceMetric; +import org.apache.hudi.common.index.vector.VectorStalePolicy; + +import java.io.Serializable; +import java.util.Objects; + +/** + * Engine-neutral vector search request (RFC-109 §1). Carries only the query intent and budget; + * no engine, storage, or SQL types. Adapters (Spark/Flink/Java) translate their inputs into this. + * + *

{@code queryInstant} pins the table snapshot for the entire request; when null the executor + * resolves the latest completed instant and records it in the result snapshot. + */ +public final class VectorSearchRequest implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String vectorColumn; + private final float[] queryVector; + private final VectorDistanceMetric metric; + private final int topK; + private final int nprobe; + private final int refineFactor; + private final boolean exactRerank; + private final VectorStalePolicy stalePolicy; + private final String queryInstant; + private final VectorSearchBudget budget; + + public VectorSearchRequest(String vectorColumn, + float[] queryVector, + VectorDistanceMetric metric, + int topK, + int nprobe, + int refineFactor, + boolean exactRerank, + String queryInstant, + VectorSearchBudget budget) { + this(vectorColumn, queryVector, metric, topK, nprobe, refineFactor, exactRerank, + VectorStalePolicy.FAIL, queryInstant, budget); + } + + public VectorSearchRequest(String vectorColumn, + float[] queryVector, + VectorDistanceMetric metric, + int topK, + int nprobe, + int refineFactor, + boolean exactRerank, + VectorStalePolicy stalePolicy, + String queryInstant, + VectorSearchBudget budget) { + this.vectorColumn = Objects.requireNonNull(vectorColumn, "vectorColumn"); + this.queryVector = Objects.requireNonNull(queryVector, "queryVector"); + this.metric = Objects.requireNonNull(metric, "metric"); + this.topK = topK; + this.nprobe = nprobe; + this.refineFactor = refineFactor; + this.exactRerank = exactRerank; + this.stalePolicy = Objects.requireNonNull(stalePolicy, "stalePolicy"); + this.queryInstant = queryInstant; + this.budget = Objects.requireNonNull(budget, "budget"); + } + + public String getVectorColumn() { + return vectorColumn; + } + + public float[] getQueryVector() { + return queryVector; + } + + public VectorDistanceMetric getMetric() { + return metric; + } + + public int getTopK() { + return topK; + } + + public int getNprobe() { + return nprobe; + } + + public int getRefineFactor() { + return refineFactor; + } + + public boolean isExactRerank() { + return exactRerank; + } + + public VectorStalePolicy getStalePolicy() { + return stalePolicy; + } + + /** Pinned table instant for the request, or null to resolve the latest completed instant. */ + public String getQueryInstant() { + return queryInstant; + } + + public VectorSearchBudget getBudget() { + return budget; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchResult.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchResult.java new file mode 100644 index 0000000000000..ea42bfa839908 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchResult.java @@ -0,0 +1,55 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import java.io.Serializable; + +/** + * One engine-neutral final result row from a vector search (RFC-109 §1): the logical record key, + * the exact metric distance (squared L2 kept internally, surfaced per the requested metric), and + * the live record location the value was read from. + */ +public final class VectorSearchResult implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String recordKey; + private final double distance; + private final HoodieRecordGlobalLocation location; + + public VectorSearchResult(String recordKey, double distance, HoodieRecordGlobalLocation location) { + this.recordKey = recordKey; + this.distance = distance; + this.location = location; + } + + public String getRecordKey() { + return recordKey; + } + + public double getDistance() { + return distance; + } + + public HoodieRecordGlobalLocation getLocation() { + return location; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchSnapshot.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchSnapshot.java new file mode 100644 index 0000000000000..0c6ef075a69ae --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchSnapshot.java @@ -0,0 +1,49 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; +import java.util.Objects; + +/** + * The single pinned snapshot used for an entire vector search: one table instant shared by the MDT + * index read, the RLI finalist lookup, file-slice resolution, and the base-table exact fetch, plus + * the resolved {@link VectorIndexSnapshot} generation identity (RFC-109 §7). Using one instant + * across all reads is what makes freshness arbitration correct. + */ +public final class VectorSearchSnapshot implements Serializable { + + private static final long serialVersionUID = 1L; + + private final String tableInstant; + private final VectorIndexSnapshot vectorIndex; + + public VectorSearchSnapshot(String tableInstant, VectorIndexSnapshot vectorIndex) { + this.tableInstant = Objects.requireNonNull(tableInstant, "tableInstant"); + this.vectorIndex = Objects.requireNonNull(vectorIndex, "vectorIndex"); + } + + public String getTableInstant() { + return tableInstant; + } + + public VectorIndexSnapshot getVectorIndex() { + return vectorIndex; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchStatus.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchStatus.java new file mode 100644 index 0000000000000..bfad3961f407b --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSearchStatus.java @@ -0,0 +1,37 @@ +/* + * 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.hudi.common.index.vector.search; + +/** + * Terminal status of a vector search request (RFC-109 §1). + * + *

    + *
  • {@code COMPLETED}: K live exact results returned.
    • + *
  • {@code PARTIAL}: fewer than K returned because the retained candidate pool was exhausted + * (not a deadline), with {@link DeadlinePolicy#RETURN_PARTIAL} in effect.
    • + *
  • {@code DEADLINE_EXCEEDED}: the request budget expired; results may be partial.
    • + *
  • {@code FAILED}: an error prevented completion.
    • + *
+ */ +public enum VectorSearchStatus { + COMPLETED, + PARTIAL, + DEADLINE_EXCEEDED, + FAILED +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSnapshotResolver.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSnapshotResolver.java new file mode 100644 index 0000000000000..8300396dc7368 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorSnapshotResolver.java @@ -0,0 +1,33 @@ +/* + * 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.hudi.common.index.vector.search; + +import java.io.Serializable; + +/** + * Resolves the single pinned {@link VectorSearchSnapshot} for a request (RFC-109 §7, §11): + * the table instant (from {@code request.queryInstant} or the latest completed instant) plus the + * active {@link VectorIndexSnapshot} generation identity. Injected so the common executor stays + * engine-neutral — engine adapters provide the metadata-backed implementation. + */ +@FunctionalInterface +public interface VectorSnapshotResolver extends Serializable { + + VectorSearchSnapshot resolve(VectorSearchRequest request); +} diff --git a/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorTopKAccumulator.java b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorTopKAccumulator.java new file mode 100644 index 0000000000000..8480dfaba3a71 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/common/index/vector/search/VectorTopKAccumulator.java @@ -0,0 +1,79 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * Accumulates the top-K live exact results during rerank/continuation (RFC-109 §10). Deduplicates + * by logical record key (keeping the smaller distance) so a record surfaced by both its posting and + * a key-fallback fetch is counted once, and reports how many live results are held so the + * continuation loop knows whether it still needs more candidates. + * + *

Distances are the order-preserving ranking distances from {@link ExactVectorScorer} (smaller = + * more similar); squared L2 is kept internally. + */ +public final class VectorTopKAccumulator { + + private final int topK; + private final Map bestByKey; + + public VectorTopKAccumulator(int topK) { + if (topK <= 0) { + throw new IllegalArgumentException("topK must be positive, got: " + topK); + } + this.topK = topK; + this.bestByKey = new HashMap<>(); + } + + /** Offer one exact-scored live record; keeps the smaller distance per record key. */ + public void offer(String recordKey, double distance, HoodieRecordGlobalLocation location) { + VectorSearchResult existing = bestByKey.get(recordKey); + if (existing == null || distance < existing.getDistance()) { + bestByKey.put(recordKey, new VectorSearchResult(recordKey, distance, location)); + } + } + + /** Number of distinct live records accumulated so far. */ + public int liveCount() { + return bestByKey.size(); + } + + /** Whether fewer than K distinct live results have been accumulated. */ + public boolean needsMore() { + return bestByKey.size() < topK; + } + + /** The current top-K live results, ascending by distance (ties broken by record key). */ + public List topK() { + List all = new ArrayList<>(bestByKey.values()); + all.sort(Comparator.comparingDouble(VectorSearchResult::getDistance) + .thenComparing(VectorSearchResult::getRecordKey)); + if (all.size() > topK) { + return new ArrayList<>(all.subList(0, topK)); + } + return all; + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieIndexVersion.java b/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieIndexVersion.java index 91b534d8d9f2e..bdafe3e3937fb 100644 --- a/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieIndexVersion.java +++ b/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieIndexVersion.java @@ -107,6 +107,9 @@ public static HoodieIndexVersion getCurrentVersion(HoodieTableVersion tableVersi } return V1; + case VECTOR_INDEX: + return V1; + case FILES: return V1; diff --git a/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataPayload.java b/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataPayload.java index 810cbecacc036..a7ae138d2b844 100644 --- a/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataPayload.java +++ b/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataPayload.java @@ -24,6 +24,15 @@ import org.apache.hudi.avro.model.HoodieMetadataRecord; import org.apache.hudi.avro.model.HoodieRecordIndexInfo; import org.apache.hudi.avro.model.HoodieSecondaryIndexInfo; +import org.apache.hudi.avro.model.HoodieVectorIndexActiveManifest; +import org.apache.hudi.avro.model.HoodieVectorIndexCentroids; +import org.apache.hudi.avro.model.HoodieVectorIndexClusterStats; +import org.apache.hudi.avro.model.HoodieVectorIndexManifest; +import org.apache.hudi.avro.model.HoodieVectorIndexPostingBlock; +import org.apache.hudi.avro.model.HoodieVectorIndexPostingDelta; +import org.apache.hudi.avro.model.HoodieVectorIndexQuantizer; +import org.apache.hudi.avro.model.HoodieVectorIndexSourceInstantMarker; +import org.apache.hudi.avro.model.HoodieVectorIndexTombstone; import org.apache.hudi.common.avro.AvroSchemaCache; import org.apache.hudi.common.fs.FSUtils; import org.apache.hudi.common.model.EmptyHoodieRecordPayload; @@ -121,6 +130,7 @@ public class HoodieMetadataPayload implements HoodieRecordPayload orderingVal) { @@ -253,6 +300,410 @@ protected HoodieMetadataPayload(String key, HoodieSecondaryIndexInfo secondaryIn this(key, MetadataPartitionType.SECONDARY_INDEX.getRecordType(), null, null, null, null, secondaryIndexMetadata, secondaryIndexMetadata.getIsDeleted()); } + protected HoodieMetadataPayload(String key, Object vectorIndexInfo) { + this.key = key; + this.type = MetadataPartitionType.VECTOR_INDEX.getRecordType(); + this.vectorIndexMetadata = vectorIndexInfo; + this.isDeletedRecord = vectorIndexInfo instanceof HoodieVectorIndexTombstone; + } + + /** + * Create the singleton reader-visible generation pointer. + */ + public static HoodieRecord createVectorIndexActiveManifestRecord( + Integer activeGeneration, String metadataPartitionPath) { + String recordKey = VectorIndexMetadataKey.activeManifest(); + HoodieVectorIndexActiveManifest manifest = new HoodieVectorIndexActiveManifest(1, activeGeneration); + return new HoodieAvroRecord<>( + new HoodieKey(recordKey, metadataPartitionPath), + new HoodieMetadataPayload(recordKey, manifest)); + } + + /** + * Create the generation-one centroid record for the given index partition. + */ + public static HoodieRecord createVectorIndexCentroidsRecord( + ByteBuffer centroidBytes, String partitionPath) { + HoodieVectorIndexCentroids centroids = new HoodieVectorIndexCentroids( + ByteBuffer.allocate(0), + centroidBytes, + ByteBuffer.allocate(0)); + String recordKey = VectorIndexMetadataKey.centroids(1, 0); + HoodieMetadataPayload payload = new HoodieMetadataPayload(recordKey, centroids); + HoodieKey key = new HoodieKey(recordKey, partitionPath); + return new HoodieAvroRecord<>(key, payload); + } + + public static HoodieRecord createVectorIndexCentroidsRecord( + int generation, + int chunk, + ByteBuffer clusterIds, + ByteBuffer centroidBytes, + ByteBuffer clusterRadii, + String partitionPath) { + String recordKey = VectorIndexMetadataKey.centroids(generation, chunk); + HoodieVectorIndexCentroids centroids = new HoodieVectorIndexCentroids( + clusterIds, centroidBytes, clusterRadii); + return new HoodieAvroRecord<>( + new HoodieKey(recordKey, partitionPath), + new HoodieMetadataPayload(recordKey, centroids)); + } + + /** + * Create the generation-one quantizer metadata record for the given index partition. + */ + public static HoodieRecord createVectorIndexQuantizerMetadataRecord( + String quantizerType, + int quantizedCodeBytes, + long randomSeed, + boolean assumeNormalized, + String partitionPath) { + return createVectorIndexQuantizerMetadataRecord( + quantizerType, + quantizedCodeBytes, + 1, + randomSeed, + assumeNormalized, + partitionPath); + } + + public static HoodieRecord createVectorIndexQuantizerMetadataRecord( + String quantizerType, + int quantizedCodeBytes, + int rabitqBits, + long randomSeed, + boolean assumeNormalized, + String partitionPath) { + return createVectorIndexQuantizerMetadataRecord(1, 0, quantizerType, randomSeed, null, partitionPath); + } + + public static HoodieRecord createVectorIndexQuantizerMetadataRecord( + int generation, + int chunk, + String quantizerType, + long randomSeed, + ByteBuffer rotationBytes, + String partitionPath) { + String recordKey = VectorIndexMetadataKey.quantizer(generation, chunk); + HoodieVectorIndexQuantizer quantizer = new HoodieVectorIndexQuantizer(quantizerType, randomSeed, rotationBytes); + return new HoodieAvroRecord<>( + new HoodieKey(recordKey, partitionPath), + new HoodieMetadataPayload(recordKey, quantizer)); + } + + public static HoodieRecord createVectorIndexManifestRecord( + int generation, + String generationOrdinalText, + String state, + int dim, + int dimPadded, + int codeRowBytes, + int bitsTotal, + int numExPlanes, + int numClusters, + int shardCount, + int fileGroupCount, + String metric, + boolean assumeNormalized, + boolean residualEncoding, + String vectorColumn, + int targetBlockBytes, + int vectorsPerBlock, + int blockFormatVersion, + int factorVersion, + double kappa, + double gMin, + double eps1Max, + double epsNRel, + int centroidChunkCount, + String centroidChecksum, + int splitLimit, + int mergeFloor, + String bootstrapInstant, + String verifiedFrontier, + long createdTs, + String metadataPartitionPath) { + String recordKey = VectorIndexMetadataKey.manifest(generation); + HoodieVectorIndexManifest manifest = new HoodieVectorIndexManifest( + 1, + generationOrdinalText, + state, + dim, + dimPadded, + codeRowBytes, + bitsTotal, + numExPlanes, + numClusters, + shardCount, + fileGroupCount, + metric, + assumeNormalized, + residualEncoding, + vectorColumn == null ? "" : vectorColumn, + targetBlockBytes, + vectorsPerBlock, + blockFormatVersion, + factorVersion, + kappa, + gMin, + eps1Max, + epsNRel, + centroidChunkCount, + centroidChecksum, + splitLimit, + mergeFloor, + bootstrapInstant, + verifiedFrontier, + createdTs); + return new HoodieAvroRecord<>( + new HoodieKey(recordKey, metadataPartitionPath), + new HoodieMetadataPayload(recordKey, manifest)); + } + + public static HoodieRecord createVectorIndexClusterManifestRecord( + int generation, + int clusterId, + int shardCount, + Collection fileGroupIds, + long vectorCount, + long lastUpdatedTs, + String metadataPartitionPath) { + return createVectorIndexClusterManifestRecord( + generation, clusterId, 0, shardCount, fileGroupIds, vectorCount, lastUpdatedTs, metadataPartitionPath); + } + + public static HoodieRecord createVectorIndexClusterManifestRecord( + int generation, + int clusterId, + int routingVersion, + int shardCount, + Collection fileGroupIds, + long vectorCount, + long lastUpdatedTs, + String metadataPartitionPath) { + String recordKey = VectorIndexMetadataKey.clusterStats(generation, clusterId); + HoodieVectorIndexClusterStats stats = new HoodieVectorIndexClusterStats( + routingVersion, + shardCount, + fileGroupIds == null ? java.util.Collections.emptyList() : fileGroupIds.stream().collect(Collectors.toList()), + vectorCount, + 0L, + 0L, + null, + lastUpdatedTs); + return new HoodieAvroRecord<>( + new HoodieKey(recordKey, metadataPartitionPath), + new HoodieMetadataPayload(recordKey, stats)); + } + + public static HoodieRecord createVectorIndexClusterStatsRecord( + int generation, + int clusterId, + long liveCount, + long deltaCount, + long tombstoneCount, + String metadataPartitionPath) { + String recordKey = VectorIndexMetadataKey.clusterStats(generation, clusterId); + HoodieVectorIndexClusterStats stats = new HoodieVectorIndexClusterStats( + 0, + 1, + java.util.Collections.emptyList(), + liveCount, + deltaCount, + tombstoneCount, + null, + 0L); + return new HoodieAvroRecord<>( + new HoodieKey(recordKey, metadataPartitionPath), + new HoodieMetadataPayload(recordKey, stats)); + } + + public static HoodieRecord createVectorIndexSourceInstantMarkerRecord( + int generation, + String dataInstant, + String metadataPartitionPath) { + String recordKey = VectorIndexMetadataKey.sourceInstantMarker(generation, dataInstant); + HoodieVectorIndexSourceInstantMarker marker = new HoodieVectorIndexSourceInstantMarker(dataInstant); + return new HoodieAvroRecord<>( + new HoodieKey(recordKey, metadataPartitionPath), + new HoodieMetadataPayload(recordKey, marker)); + } + + public static HoodieRecord createVectorIndexPostingRecord( + int generation, + String recordKey, + int clusterId, + String fileGroupId, + String dataPartitionPath, + String baseInstantTime, + byte[] binaryCode, + Float scalar, + long lastUpdatedTs, + String metadataPartitionPath) { + return createVectorIndexPostingRecord( + generation, + recordKey, + clusterId, + 0, + fileGroupId, + dataPartitionPath, + baseInstantTime, + binaryCode, + null, + scalar, + null, + null, + lastUpdatedTs, + metadataPartitionPath); + } + + public static HoodieRecord createVectorIndexPostingRecord( + int generation, + String recordKey, + int clusterId, + int shardId, + String fileGroupId, + String dataPartitionPath, + String baseInstantTime, + byte[] binaryCode, + Float scalar, + long lastUpdatedTs, + String metadataPartitionPath) { + return createVectorIndexPostingRecord( + generation, + recordKey, + clusterId, + shardId, + fileGroupId, + dataPartitionPath, + baseInstantTime, + binaryCode, + null, + scalar, + null, + null, + lastUpdatedTs, + metadataPartitionPath); + } + + public static HoodieRecord createVectorIndexPostingRecord( + int generation, + String recordKey, + int clusterId, + int shardId, + String fileGroupId, + String dataPartitionPath, + String baseInstantTime, + byte[] binaryCode, + byte[] extendedCode, + Float scalar, + Float additiveFactor, + Float rescaleFactor, + long lastUpdatedTs, + String metadataPartitionPath) { + return createVectorIndexPostingRecord( + generation, + recordKey, + clusterId, + shardId, + fileGroupId, + dataPartitionPath, + baseInstantTime, + binaryCode, + extendedCode, + scalar, + additiveFactor, + rescaleFactor, + 0.0f, + 0.0f, + 0.0f, + lastUpdatedTs, + metadataPartitionPath); + } + + public static HoodieRecord createVectorIndexPostingRecord( + int generation, + String recordKey, + int clusterId, + int shardId, + String fileGroupId, + String dataPartitionPath, + String baseInstantTime, + byte[] binaryCode, + byte[] extendedCode, + Float scalar, + Float additiveFactor, + Float rescaleFactor, + Float additiveFactor1, + Float rescaleFactor1, + Float error1, + long lastUpdatedTs, + String metadataPartitionPath) { + String metadataRecordKey = VectorIndexMetadataKey.postingDelta(generation, clusterId, shardId, recordKey); + byte[] code = mergeCodeRows(binaryCode, extendedCode); + HoodieVectorIndexPostingDelta delta = new HoodieVectorIndexPostingDelta( + recordKey, + ByteBuffer.wrap(code), + additiveFactor1 == null ? 0.0f : additiveFactor1, + rescaleFactor1 == null ? 0.0f : rescaleFactor1, + error1 == null ? 0.0f : error1, + additiveFactor == null ? 0.0f : additiveFactor, + rescaleFactor == null ? 0.0f : rescaleFactor, + scalar == null ? 0.0f : scalar, + null, + fileGroupId == null ? "" : fileGroupId, + dataPartitionPath == null ? "" : dataPartitionPath, + baseInstantTime == null ? "" : baseInstantTime, + -1L); + HoodieMetadataPayload payload = new HoodieMetadataPayload(metadataRecordKey, delta); + HoodieKey key = new HoodieKey(metadataRecordKey, metadataPartitionPath); + return new HoodieAvroRecord<>(key, payload); + } + + public static HoodieRecord createVectorIndexPostingBlockRecord( + int generation, + int clusterId, + int shardId, + long blockId, + HoodieVectorIndexPostingBlock postingBlock, + String metadataPartitionPath) { + String metadataRecordKey = VectorIndexMetadataKey.postingBlock(generation, clusterId, shardId, blockId); + HoodieMetadataPayload payload = new HoodieMetadataPayload(metadataRecordKey, postingBlock); + HoodieKey key = new HoodieKey(metadataRecordKey, metadataPartitionPath); + return new HoodieAvroRecord<>(key, payload); + } + + /** + * Create a tombstone for a canonical vector posting record. + */ + public static HoodieRecord createVectorIndexPostingDeleteRecord( + int generation, + String recordKey, + int clusterId, + int shardId, + String metadataPartitionPath) { + String metadataRecordKey = VectorIndexMetadataKey.postingDelta(generation, clusterId, shardId, recordKey); + HoodieVectorIndexTombstone tombstone = new HoodieVectorIndexTombstone("", "MANUAL"); + HoodieMetadataPayload payload = new HoodieMetadataPayload(metadataRecordKey, tombstone); + HoodieKey key = new HoodieKey(metadataRecordKey, metadataPartitionPath); + return new HoodieAvroRecord<>(key, payload); + } + + /** + * Return the vector index metadata if present. + */ + public Option getVectorIndexMetadata() { + return Option.ofNullable(vectorIndexMetadata); + } + + private static byte[] mergeCodeRows(byte[] binaryCode, byte[] extendedCode) { + byte[] first = binaryCode == null ? new byte[0] : binaryCode; + byte[] second = extendedCode == null ? new byte[0] : extendedCode; + byte[] merged = Arrays.copyOf(first, first.length + second.length); + System.arraycopy(second, 0, merged, first.length, second.length); + return merged; + } + protected HoodieMetadataPayload(String key, int type, Map filesystemMetadata, HoodieMetadataBloomFilter metadataBloomFilter, @@ -429,7 +880,7 @@ public Option getInsertValue(Schema schema, Properties properties if (schema == null || schema == HOODIE_METADATA_AVRO_SCHEMA) { // If the schema is same or none is provided, we can return the record directly HoodieMetadataRecord record = new HoodieMetadataRecord(key, type, filesystemMetadata, bloomFilterMetadata, - columnStatMetadata, recordIndexMetadata, secondaryIndexMetadata); + columnStatMetadata, recordIndexMetadata, secondaryIndexMetadata, vectorIndexMetadata); return Option.of(record); } else { // Otherwise, the assumption is that the schema required contains the metadata fields so we construct a new GenericRecord with these fields @@ -451,6 +902,9 @@ public Option getInsertValue(Schema schema, Properties properties if (secondaryIndexMetadata != null) { record.put(SECONDARY_INDEX_METADATA_FIELD_OFFSET, secondaryIndexMetadata); } + if (vectorIndexMetadata != null) { + record.put(VECTOR_INDEX_METADATA_FIELD_OFFSET, vectorIndexMetadata); + } return Option.of(record); } } diff --git a/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieTableMetadataUtil.java b/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieTableMetadataUtil.java index 202a2e6a0b330..e682518670b95 100644 --- a/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieTableMetadataUtil.java +++ b/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieTableMetadataUtil.java @@ -202,6 +202,8 @@ public class HoodieTableMetadataUtil { public static final String PARTITION_NAME_EXPRESSION_INDEX_PREFIX = "expr_index_"; public static final String PARTITION_NAME_SECONDARY_INDEX = "secondary_index"; public static final String PARTITION_NAME_SECONDARY_INDEX_PREFIX = "secondary_index_"; + public static final String PARTITION_NAME_VECTOR_INDEX = "vector_index"; + public static final String PARTITION_NAME_VECTOR_INDEX_PREFIX = "vector_index_"; // Average size of a record saved within the record index. // Record index has a fixed size schema. This has been calculated based on experiments with default settings @@ -213,6 +215,71 @@ public class HoodieTableMetadataUtil { HoodieRecord.HoodieMetadataField.PARTITION_PATH_METADATA_FIELD.getFieldName(), HoodieRecord.HoodieMetadataField.COMMIT_TIME_METADATA_FIELD.getFieldName())); + public static String getVectorIndexGenerationManifestKey(int generationId) { + return VectorIndexMetadataKey.manifest(generationId); + } + + public static boolean isVectorIndexGenerationManifestKey(String recordKey) { + return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_MANIFEST); + } + + public static boolean isVectorIndexQuantizerKey(String recordKey) { + return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_QUANTIZER); + } + + public static boolean isVectorIndexCentroidsKey(String recordKey) { + return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_CENTROIDS); + } + + public static String getVectorIndexClusterKey(int generationId, int clusterId) { + return VectorIndexMetadataKey.clusterStats(generationId, clusterId); + } + + public static boolean isVectorIndexClusterKey(String recordKey) { + return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_CLUSTER_STATS); + } + + public static String getVectorIndexPostingKey(int generationId, int clusterId, int shardId, String recordKey) { + return VectorIndexMetadataKey.postingDelta(generationId, clusterId, shardId, recordKey); + } + + public static String getVectorIndexPostingPrefix(int generationId, int clusterId) { + return VectorIndexMetadataKey.postingPrefix(generationId, clusterId, 0).substring(0, 9); + } + + public static int toVectorGenerationId(String generationId) { + try { + int parsed = Integer.parseInt(generationId); + if (parsed <= 0) { + throw new IllegalArgumentException( + "Vector generation id must be a positive ordinal, got: " + generationId); + } + return parsed; + } catch (NumberFormatException e) { + throw new IllegalArgumentException( + "Vector generation id must be a positive integer ordinal, got: " + generationId, e); + } + } + + public static String getVectorIndexPostingPrefix(int generationId, int clusterId, int shardId) { + return VectorIndexMetadataKey.postingPrefix(generationId, clusterId, shardId); + } + + public static boolean isVectorIndexPostingKey(String recordKey) { + return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_POSTING); + } + + private static boolean hasVectorIndexFamily(String recordKey, int family) { + return recordKey != null && !recordKey.isEmpty() + && Byte.toUnsignedInt(VectorIndexMetadataKey.decode(recordKey)[0]) == family; + } + + public static String getVectorIndexPostingRecordKey(String metadataRecordKey) { + return isVectorIndexPostingKey(metadataRecordKey) + ? VectorIndexMetadataKey.postingRecordKey(metadataRecordKey) + : null; + } + // The maximum allowed precision and scale as per the payload schema. See DecimalWrapper in HoodieMetadata.avsc: // https://github.com/apache/hudi/blob/45dedd819e56e521148bde51a3dfa4e472ea70cd/hudi-common/src/main/avro/HoodieMetadata.avsc#L247 private static final int DECIMAL_MAX_PRECISION = 30; @@ -960,6 +1027,39 @@ public static int mapRecordKeyToFileGroupIndex(String recordKey, int numFileGrou return Math.abs(Math.abs(h) % numFileGroups); } + public static int mapVectorPostingKeyToFileGroupIndex(String recordKey, int numFileGroups) { + if (numFileGroups <= 1) { + return 0; + } + // Route by the IVF cluster id so that all postings of a cluster (every shard, every + // record key) land in the same MDT file group. With fileGroupCount == num_clusters this + // yields exactly one cluster per file group; with fewer file groups, clusters are spread + // round-robin (clusterId % numFileGroups) while a single cluster still stays whole. + int clusterId = extractVectorClusterId(recordKey); + if (clusterId < 0) { + // Not a cluster-bearing posting key (e.g. manifest/centroid rows): fall back to stable + // string hashing on the routing key prefix. + return mapRecordKeyToFileGroupIndex(getVectorIndexPostingRoutingKey(recordKey), numFileGroups); + } + return Math.floorMod(clusterId, numFileGroups); + } + + /** + * Extract the IVF cluster id from a vector posting key/prefix of the form + * {@code P|||...}. Returns {@code -1} when {@code recordKey} is not a + * vector posting key or the cluster segment cannot be parsed. + */ + public static int extractVectorClusterId(String recordKey) { + return isVectorIndexPostingKey(recordKey) ? VectorIndexMetadataKey.postingClusterId(recordKey) : -1; + } + + public static String getVectorIndexPostingRoutingKey(String metadataRecordKey) { + if (!isVectorIndexPostingKey(metadataRecordKey)) { + return metadataRecordKey; + } + return metadataRecordKey.length() >= 9 ? metadataRecordKey.substring(0, 9) : metadataRecordKey; + } + /** * Get the latest file slices for a Metadata Table partition. If the file slice is * because of pending compaction instant, then merge the file slice with the one diff --git a/hudi-common/src/main/java/org/apache/hudi/metadata/MetadataPartitionType.java b/hudi-common/src/main/java/org/apache/hudi/metadata/MetadataPartitionType.java index 6311a7cf15804..b4a386d12585f 100644 --- a/hudi-common/src/main/java/org/apache/hudi/metadata/MetadataPartitionType.java +++ b/hudi-common/src/main/java/org/apache/hudi/metadata/MetadataPartitionType.java @@ -77,6 +77,7 @@ import static org.apache.hudi.metadata.HoodieMetadataPayload.SCHEMA_FIELD_ID_COLUMN_STATS; import static org.apache.hudi.metadata.HoodieMetadataPayload.SCHEMA_FIELD_ID_RECORD_INDEX; import static org.apache.hudi.metadata.HoodieMetadataPayload.SCHEMA_FIELD_ID_SECONDARY_INDEX; +import static org.apache.hudi.metadata.HoodieMetadataPayload.SCHEMA_FIELD_ID_VECTOR_INDEX; import static org.apache.hudi.metadata.HoodieMetadataPayload.SCHEMA_FIELD_NAME_METADATA; import static org.apache.hudi.metadata.HoodieMetadataPayload.SECONDARY_INDEX_FIELD_IS_DELETED; import static org.apache.hudi.metadata.HoodieTableMetadataUtil.PARTITION_NAME_EXPRESSION_INDEX; @@ -244,6 +245,42 @@ public SerializableBiFunction getFileGroupMappingFunct return HoodieTableMetadataUtil.getSecondaryKeyToFileGroupMappingFunction(indexVersion.greaterThanOrEquals(HoodieIndexVersion.V2)); } }, + VECTOR_INDEX(HoodieTableMetadataUtil.PARTITION_NAME_VECTOR_INDEX_PREFIX, "vector-index-", 8) { + @Override + public boolean isMetadataPartitionEnabled(HoodieMetadataConfig metadataConfig, HoodieTableConfig tableConfig) { + // Vector index is enabled explicitly via CREATE INDEX, not via metadata config flags. + return false; + } + + @Override + public boolean isMetadataPartitionAvailable(HoodieTableMetaClient metaClient) { + if (metaClient.getIndexMetadata().isPresent()) { + return metaClient.getIndexMetadata().get().getIndexDefinitions().values().stream() + .anyMatch(indexDef -> indexDef.getIndexName().startsWith(HoodieTableMetadataUtil.PARTITION_NAME_VECTOR_INDEX_PREFIX)); + } + return false; + } + + @Override + public void constructMetadataPayload(HoodieMetadataPayload payload, GenericRecord record) { + GenericRecord vectorIndexRecord = getNestedFieldValue(record, SCHEMA_FIELD_ID_VECTOR_INDEX); + checkState(vectorIndexRecord != null, + "Valid VectorIndexMetadata record expected for type: " + MetadataPartitionType.VECTOR_INDEX.getRecordType()); + payload.vectorIndexMetadata = vectorIndexRecord; + } + + @Override + public String getPartitionPath(HoodieTableMetaClient metaClient, String indexName) { + return metaClient.getIndexForMetadataPartition(indexName) + .map(HoodieIndexDefinition::getIndexName) + .orElseThrow(() -> new IllegalArgumentException("Index definition is not present for index: " + indexName)); + } + + @Override + public SerializableBiFunction getFileGroupMappingFunction(HoodieIndexVersion indexVersion) { + return HoodieTableMetadataUtil::mapVectorPostingKeyToFileGroupIndex; + } + }, PARTITION_STATS(HoodieTableMetadataUtil.PARTITION_NAME_PARTITION_STATS, "partition-stats-", 6) { @Override public boolean isMetadataPartitionEnabled(HoodieMetadataConfig metadataConfig, HoodieTableConfig tableConfig) { @@ -468,6 +505,7 @@ public static MetadataPartitionType[] getValidValues(HoodieTableVersion tableVer .stream() .filter(type -> type != SECONDARY_INDEX && type != EXPRESSION_INDEX + && type != VECTOR_INDEX && type != PARTITION_STATS) .toArray(MetadataPartitionType[]::new); } diff --git a/hudi-common/src/main/java/org/apache/hudi/metadata/VectorClusterRawKey.java b/hudi-common/src/main/java/org/apache/hudi/metadata/VectorClusterRawKey.java new file mode 100644 index 0000000000000..4e6c22f9f982b --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/metadata/VectorClusterRawKey.java @@ -0,0 +1,37 @@ +/* + * 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.hudi.metadata; + +import lombok.Value; + +/** + * Raw key for vector cluster metadata records. + */ +@Value +public class VectorClusterRawKey implements RawKey { + + int generationId; + int clusterId; + + @Override + public String encode() { + return VectorIndexMetadataKey.clusterStats(generationId, clusterId); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/metadata/VectorGenerationManifestRawKey.java b/hudi-common/src/main/java/org/apache/hudi/metadata/VectorGenerationManifestRawKey.java new file mode 100644 index 0000000000000..57289a5901d91 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/metadata/VectorGenerationManifestRawKey.java @@ -0,0 +1,36 @@ +/* + * 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.hudi.metadata; + +import lombok.Value; + +/** + * Raw key for vector generation manifest records. + */ +@Value +public class VectorGenerationManifestRawKey implements RawKey { + + int generationId; + + @Override + public String encode() { + return VectorIndexMetadataKey.manifest(generationId); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/metadata/VectorIndexMetadataKey.java b/hudi-common/src/main/java/org/apache/hudi/metadata/VectorIndexMetadataKey.java new file mode 100644 index 0000000000000..bb5072f3f742a --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/metadata/VectorIndexMetadataKey.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.hudi.metadata; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; + +/** + * Binary-sortable key contract for records in the vector-index metadata partition. + * + *

Hudi metadata record keys are strings today, so the fixed binary prefix is represented as + * ISO-8859-1 characters. That preserves the unsigned byte value of each component and keeps the + * HFile key bytes lexicographically sortable by family, generation, cluster, shard, and block. + */ +public final class VectorIndexMetadataKey { + + public static final int FAMILY_ACTIVE_MANIFEST = 0x00; + public static final int FAMILY_MANIFEST = 0x01; + public static final int FAMILY_QUANTIZER = 0x02; + public static final int FAMILY_CENTROIDS = 0x03; + public static final int FAMILY_CLUSTER_STATS = 0x04; + public static final int FAMILY_POSTING = 0x10; + public static final int FAMILY_SOURCE_INSTANT_MARKER = 0x20; + + public static final long MAX_PACKED_BLOCK_ID = 0xFFFDFFFFL; + public static final long FIRST_RESERVED_BLOCK_ID = 0xFFFE0000L; + public static final long LAST_RESERVED_BLOCK_ID = 0xFFFEFFFFL; + public static final long DELTA_BLOCK_ID = 0xFFFFFFFFL; + + private VectorIndexMetadataKey() { + } + + public static String activeManifest() { + return encode(ByteBuffer.allocate(1).put((byte) FAMILY_ACTIVE_MANIFEST)); + } + + public static String manifest(int generation) { + return encode(putUnsignedInt(ByteBuffer.allocate(5).put((byte) FAMILY_MANIFEST), Integer.toUnsignedLong(generation))); + } + + public static String quantizer(int generation, int chunk) { + ByteBuffer buffer = ByteBuffer.allocate(7).order(ByteOrder.BIG_ENDIAN); + buffer.put((byte) FAMILY_QUANTIZER); + putUnsignedInt(buffer, Integer.toUnsignedLong(generation)); + putUnsignedShort(buffer, chunk); + return encode(buffer); + } + + public static String centroids(int generation, int chunk) { + ByteBuffer buffer = ByteBuffer.allocate(7).order(ByteOrder.BIG_ENDIAN); + buffer.put((byte) FAMILY_CENTROIDS); + putUnsignedInt(buffer, Integer.toUnsignedLong(generation)); + putUnsignedShort(buffer, chunk); + return encode(buffer); + } + + public static String clusterStats(int generation, int clusterId) { + ByteBuffer buffer = ByteBuffer.allocate(9).order(ByteOrder.BIG_ENDIAN); + buffer.put((byte) FAMILY_CLUSTER_STATS); + putUnsignedInt(buffer, Integer.toUnsignedLong(generation)); + putUnsignedInt(buffer, Integer.toUnsignedLong(clusterId)); + return encode(buffer); + } + + public static String postingBlock(int generation, int clusterId, int shard, long blockId) { + checkUnsignedInt(blockId, "blockId"); + ByteBuffer buffer = postingPrefixBuffer(generation, clusterId, shard, 4); + putUnsignedInt(buffer, blockId); + return encode(buffer); + } + + public static String postingDelta(int generation, int clusterId, int shard, String recordKey) { + byte[] recordKeyBytes = recordKey.getBytes(StandardCharsets.UTF_8); + ByteBuffer buffer = postingPrefixBuffer(generation, clusterId, shard, 4 + recordKeyBytes.length); + putUnsignedInt(buffer, DELTA_BLOCK_ID); + buffer.put(recordKeyBytes); + return encode(buffer); + } + + public static String postingPrefix(int generation, int clusterId, int shard) { + return encode(postingPrefixBuffer(generation, clusterId, shard, 0)); + } + + public static String sourceInstantMarker(int generation, String dataInstant) { + if (dataInstant == null || dataInstant.isEmpty()) { + throw new IllegalArgumentException("dataInstant must not be empty"); + } + byte[] instantBytes = dataInstant.getBytes(StandardCharsets.UTF_8); + ByteBuffer buffer = ByteBuffer.allocate(5 + instantBytes.length).order(ByteOrder.BIG_ENDIAN); + buffer.put((byte) FAMILY_SOURCE_INSTANT_MARKER); + putUnsignedInt(buffer, Integer.toUnsignedLong(generation)); + buffer.put(instantBytes); + return encode(buffer); + } + + public static String sourceInstantMarkerPrefix(int generation) { + ByteBuffer buffer = ByteBuffer.allocate(5).order(ByteOrder.BIG_ENDIAN); + buffer.put((byte) FAMILY_SOURCE_INSTANT_MARKER); + putUnsignedInt(buffer, Integer.toUnsignedLong(generation)); + return encode(buffer); + } + + public static String sourceInstant(String key) { + byte[] bytes = decode(key); + if (bytes.length <= 5 || Byte.toUnsignedInt(bytes[0]) != FAMILY_SOURCE_INSTANT_MARKER) { + return null; + } + return new String(bytes, 5, bytes.length - 5, StandardCharsets.UTF_8); + } + + public static String exclusiveEnd(String prefix) { + byte[] bytes = decode(prefix); + for (int i = bytes.length - 1; i >= 0; i--) { + int value = Byte.toUnsignedInt(bytes[i]); + if (value != 0xFF) { + bytes[i] = (byte) (value + 1); + return encode(Arrays.copyOf(bytes, i + 1)); + } + } + return null; + } + + public static byte[] decode(String key) { + return key.getBytes(StandardCharsets.ISO_8859_1); + } + + public static int postingClusterId(String key) { + byte[] bytes = decode(key); + if (bytes.length < 9 || Byte.toUnsignedInt(bytes[0]) != FAMILY_POSTING) { + return -1; + } + return readInt(bytes, 5); + } + + public static int postingShard(String key) { + byte[] bytes = decode(key); + if (bytes.length < 11 || Byte.toUnsignedInt(bytes[0]) != FAMILY_POSTING) { + return -1; + } + return (Byte.toUnsignedInt(bytes[9]) << 8) | Byte.toUnsignedInt(bytes[10]); + } + + public static String postingRecordKey(String key) { + byte[] bytes = decode(key); + if (bytes.length <= 15 || Byte.toUnsignedInt(bytes[0]) != FAMILY_POSTING + || Integer.toUnsignedLong(readInt(bytes, 11)) != DELTA_BLOCK_ID) { + return null; + } + return new String(bytes, 15, bytes.length - 15, StandardCharsets.UTF_8); + } + + static int compareUnsigned(String left, String right) { + byte[] leftBytes = decode(left); + byte[] rightBytes = decode(right); + int length = Math.min(leftBytes.length, rightBytes.length); + for (int i = 0; i < length; i++) { + int comparison = Integer.compare(Byte.toUnsignedInt(leftBytes[i]), Byte.toUnsignedInt(rightBytes[i])); + if (comparison != 0) { + return comparison; + } + } + return Integer.compare(leftBytes.length, rightBytes.length); + } + + private static int readInt(byte[] bytes, int offset) { + return (Byte.toUnsignedInt(bytes[offset]) << 24) + | (Byte.toUnsignedInt(bytes[offset + 1]) << 16) + | (Byte.toUnsignedInt(bytes[offset + 2]) << 8) + | Byte.toUnsignedInt(bytes[offset + 3]); + } + + private static ByteBuffer postingPrefixBuffer(int generation, int clusterId, int shard, int suffixBytes) { + ByteBuffer buffer = ByteBuffer.allocate(11 + suffixBytes).order(ByteOrder.BIG_ENDIAN); + buffer.put((byte) FAMILY_POSTING); + putUnsignedInt(buffer, Integer.toUnsignedLong(generation)); + putUnsignedInt(buffer, Integer.toUnsignedLong(clusterId)); + putUnsignedShort(buffer, shard); + return buffer; + } + + private static ByteBuffer putUnsignedInt(ByteBuffer buffer, long value) { + checkUnsignedInt(value, "value"); + buffer.putInt((int) value); + return buffer; + } + + private static void putUnsignedShort(ByteBuffer buffer, int value) { + if (value < 0 || value > 0xFFFF) { + throw new IllegalArgumentException("value must fit in unsigned short: " + value); + } + buffer.putShort((short) value); + } + + private static void checkUnsignedInt(long value, String name) { + if (value < 0 || value > 0xFFFFFFFFL) { + throw new IllegalArgumentException(name + " must fit in unsigned int: " + value); + } + } + + private static String encode(ByteBuffer buffer) { + return new String(buffer.array(), StandardCharsets.ISO_8859_1); + } + + private static String encode(byte[] bytes) { + return new String(bytes, StandardCharsets.ISO_8859_1); + } +} diff --git a/hudi-common/src/main/java/org/apache/hudi/metadata/VectorPostingPrefixRawKey.java b/hudi-common/src/main/java/org/apache/hudi/metadata/VectorPostingPrefixRawKey.java new file mode 100644 index 0000000000000..8c8ac2233a2a7 --- /dev/null +++ b/hudi-common/src/main/java/org/apache/hudi/metadata/VectorPostingPrefixRawKey.java @@ -0,0 +1,40 @@ +/* + * 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.hudi.metadata; + +import lombok.Value; + +/** + * Raw key prefix for vector posting scans. + */ +@Value +public class VectorPostingPrefixRawKey implements RawKey { + + int generationId; + int clusterId; + Integer shardId; + + @Override + public String encode() { + return shardId == null + ? VectorIndexMetadataKey.postingPrefix(generationId, clusterId, 0).substring(0, 9) + : VectorIndexMetadataKey.postingPrefix(generationId, clusterId, shardId); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/avro/TestAvroRecordSizeEstimator.java b/hudi-common/src/test/java/org/apache/hudi/common/avro/TestAvroRecordSizeEstimator.java index 613dc13d9afb0..c29d150352527 100644 --- a/hudi-common/src/test/java/org/apache/hudi/common/avro/TestAvroRecordSizeEstimator.java +++ b/hudi-common/src/test/java/org/apache/hudi/common/avro/TestAvroRecordSizeEstimator.java @@ -49,7 +49,7 @@ void testEstimatingRecord() throws IOException { Assertions.assertTrue(size < 400 && size > 0); // testing generated IndexedRecord - HoodieMetadataRecord metadataRecord = new HoodieMetadataRecord("__all_partitions__", 1, new HashMap<>(), null, null, null, null); + HoodieMetadataRecord metadataRecord = new HoodieMetadataRecord("__all_partitions__", 1, new HashMap<>(), null, null, null, null, null); bufferedRecord = new BufferedRecord<>("__all_partitions__", 0, metadataRecord, 1, null); size = estimator.sizeEstimate(bufferedRecord); // size can be various for different OS / JVM version diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/MetricEstimatorIdentityTest.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/MetricEstimatorIdentityTest.java new file mode 100644 index 0000000000000..66cb36a12134b --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/MetricEstimatorIdentityTest.java @@ -0,0 +1,321 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; + +import java.util.Random; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * The metric certification suite. Three families: + * + *

IDENTITY: for random (x, c, q), the query-state composition over the NEUTRAL factors + * must equal the exact estimator composed in double precision, for both passes and every + * metric. This pins encoder and scorer to each other; every silent-recall bug this index + * has had (zero factors, folded conventions, sign drift) fails this test. + * + *

BOUND: optimisticRankingDistance must never exceed the estimator-composed distance, and the + * true distance must lie within the propagated bound at the calibrated rate. The direction + * (rip + errAbs for ALL metrics) is the designated sign-bug site. + * + *

DEGENERATE: the normative special cases compose to EXACT values. + * + *

Codes here are synthetic centered codes built directly from the residual (sign code + * c1_i = +-0.5; multibit codeEx from B-bit grid) so the suite tests the FACTOR/COMPOSITION + * contract independent of bit-packing; the golden round-trip test (separate) certifies + * packing. Together they certify the full path. + */ +public class MetricEstimatorIdentityTest { + + private static final int DIM = 96; + private static final int TRIALS = 200; + private static final int BITS = 4; + private static final double REL_TOL = 1e-5; + + @ParameterizedTest + @EnumSource(value = VectorDistanceMetric.class) + public void identityBothPasses(VectorDistanceMetric metric) { + Random rnd = new Random(42); + for (int t = 0; t < TRIALS; t++) { + Fixture f = Fixture.random(rnd, DIM, BITS, metric == VectorDistanceMetric.COSINE); + + MetricQueryState state = state(metric, f.q, false); + MetricQueryState.ClusterQuery cq = state.forRotatedCentroid(f.c); + + // ---- pass 1 ---- + double rip1 = (double) f.factors.fRescale1 * dot(f.residualQuery(), f.code1); + double got1 = state.rankingDistance(rip1, f.factors.centerRip, f.factors.residualNorm, + f.factors.vectorNorm, cq); + double exact1 = exactComposition(metric, f, dot(f.residualQuery(), f.code1), f.ipResidual1, cq); + assertRel(exact1, got1, "pass1 " + metric + " trial " + t); + + // ---- pass 2 ---- + double ripEx = (double) f.factors.fRescaleEx * dot(f.residualQuery(), f.codeEx); + double gotEx = state.rankingDistance(ripEx, f.factors.centerRip, f.factors.residualNorm, + f.factors.vectorNorm, cq); + double exactEx = exactComposition(metric, f, dot(f.residualQuery(), f.codeEx), f.ipResidualEx, cq); + assertRel(exactEx, gotEx, "pass2 " + metric + " trial " + t); + } + } + + /** + * Exact estimator composition in double precision, straight from definitions: + * rip* = n^2 * <q, code> / <r, code>; centroid component exact. + */ + private static double exactComposition(VectorDistanceMetric metric, Fixture f, + double wDotCode, double ipResidual, + MetricQueryState.ClusterQuery cq) { + double n = norm(f.r); + double rip = Math.abs(ipResidual) < 1.0e-12 + ? 0.0 : n * n * wDotCode / ipResidual; + double centerRip = dot(f.c, f.r); + switch (metric) { + case L2: + return cq.wNormSq + n * n - 2.0 * rip; + case DOT_PRODUCT: + return -(dot(f.q, f.c) + centerRip + rip); + case COSINE: + return 1.0 - (dot(f.q, f.c) + centerRip + rip) / (norm(f.q) * norm(f.x)); + default: + throw new AssertionError(); + } + } + + @ParameterizedTest + @EnumSource(value = VectorDistanceMetric.class) + public void optimisticBoundNeverExceedsEstimate(VectorDistanceMetric metric) { + Random rnd = new Random(7); + for (int t = 0; t < TRIALS; t++) { + Fixture f = Fixture.random(rnd, DIM, BITS, metric == VectorDistanceMetric.COSINE); + MetricQueryState state = state(metric, f.q, false); + MetricQueryState.ClusterQuery cq = state.forRotatedCentroid(f.c); + + double rip1 = (double) f.factors.fRescale1 * dot(f.residualQuery(), f.code1); + double est = state.rankingDistance(rip1, f.factors.centerRip, f.factors.residualNorm, + f.factors.vectorNorm, cq); + double opt = state.optimisticRankingDistance(rip1, f.factors.err1, f.factors.centerRip, + f.factors.residualNorm, f.factors.vectorNorm, cq); + // The sign-bug canary: optimistic must be <= estimate for EVERY metric + // (larger rip => smaller distance under all three finishes). + assertTrue(opt <= est + 1e-9, + metric + " optimistic bound above estimate at trial " + t + ": " + opt + " > " + est); + } + } + + @ParameterizedTest + @EnumSource(value = VectorDistanceMetric.class) + public void trueDistanceWithinBound(VectorDistanceMetric metric) { + Random rnd = new Random(11); + int violations = 0; + int trials = 2000; + for (int t = 0; t < trials; t++) { + Fixture f = Fixture.random(rnd, DIM, BITS, metric == VectorDistanceMetric.COSINE); + MetricQueryState state = state(metric, f.q, false); + MetricQueryState.ClusterQuery cq = state.forRotatedCentroid(f.c); + + double rip1 = (double) f.factors.fRescale1 * dot(f.residualQuery(), f.code1); + double opt = state.optimisticRankingDistance(rip1, f.factors.err1, f.factors.centerRip, + f.factors.residualNorm, f.factors.vectorNorm, cq); + double trueDist = trueDistance(metric, f); + if (opt > trueDist + 1e-9) { + violations++; + } + } + // kappa = 1.9 targets ~5% one-sided miss rate; allow slack for finite trials/dim. + assertTrue(violations < trials * 0.08, + metric + " bound violation rate " + violations + "/" + trials); + } + + @ParameterizedTest + @EnumSource(value = VectorDistanceMetric.class) + public void degenerateVectorAtCentroidIsExact(VectorDistanceMetric metric) { + Random rnd = new Random(3); + Fixture f = Fixture.atCentroid(rnd, DIM); + MetricQueryState state = state(metric, f.q, false); + MetricQueryState.ClusterQuery cq = state.forRotatedCentroid(f.c); + + double got = state.rankingDistance(0.0 /* rip: fRescale=0 */, f.factors.centerRip, + f.factors.residualNorm, f.factors.vectorNorm, cq); + assertRel(trueDistance(metric, f), got, "degenerate " + metric); + assertEquals(0f, f.factors.err1, 0f, "degenerate err1 must be zero (exact => prunable)"); + } + + @Test + public void l2QueryAtCentroidIsExact() { + Random rnd = new Random(5); + Fixture f = Fixture.random(rnd, DIM, BITS, false); + MetricQueryState state = state(VectorDistanceMetric.L2, f.c.clone(), false); + MetricQueryState.ClusterQuery cq = state.forRotatedCentroid(f.c); + assertTrue(cq.queryAtCentroid); + double got = state.rankingDistance(123.456 /* rip must be ignored */, f.factors.centerRip, + f.factors.residualNorm, f.factors.vectorNorm, cq); + assertRel((double) f.factors.residualNorm * f.factors.residualNorm, got, "q==c exact d2=n2"); + } + + @Test + public void zeroQueryRejectedForDotAndCosine() { + float[] zero = new float[DIM]; + assertThrows(IllegalArgumentException.class, + () -> state(VectorDistanceMetric.DOT_PRODUCT, zero, false)); + assertThrows(IllegalArgumentException.class, + () -> state(VectorDistanceMetric.COSINE, zero, false)); + state(VectorDistanceMetric.L2, zero, false); // permitted + } + + // ====================================================================================== + + /** Synthetic fixture: x, c, q, residual, both codes, neutral factors. Rotation = identity + * (a rotation is an isometry; factor/composition algebra is rotation-invariant, and the + * rotation itself is certified by the golden round-trip test). */ + private static final class Fixture { + final float[] x; + final float[] c; + final float[] q; + final float[] r; + final float[] code1; + final float[] codeEx; + final double ipResidual1; + final double ipResidualEx; + final RaBitQNeutralFactors.Factors factors; + + private Fixture(float[] x, float[] c, float[] q, float[] r, + float[] code1, float[] codeEx) { + this.x = x; + this.c = c; + this.q = q; + this.r = r; + this.code1 = code1; + this.codeEx = codeEx; + this.ipResidual1 = dot(r, code1); + this.ipResidualEx = dot(r, codeEx); + this.factors = RaBitQNeutralFactors.compute(r, c, x, ipResidual1, ipResidualEx, r.length, + RaBitQFactorConfig.defaults()); + } + + static Fixture random(Random rnd, int dim, int bits, boolean normalizeX) { + float[] c = gaussian(rnd, dim, 0.5f); + float[] x = new float[dim]; + float[] rr = gaussian(rnd, dim, 1.0f); + for (int i = 0; i < dim; i++) { + x[i] = c[i] + rr[i]; + } + if (normalizeX) { + float inv = (float) (1.0 / norm(x)); + for (int i = 0; i < dim; i++) { + x[i] *= inv; + } + } + float[] q = gaussian(rnd, dim, 1.0f); + float[] r = new float[dim]; + for (int i = 0; i < dim; i++) { + r[i] = x[i] - c[i]; + } + return new Fixture(x, c, q, r, signCode(r), multibitCode(r, bits)); + } + + static Fixture atCentroid(Random rnd, int dim) { + float[] c = gaussian(rnd, dim, 0.5f); + float[] x = c.clone(); + float[] q = gaussian(rnd, dim, 1.0f); + float[] r = new float[dim]; + return new Fixture(x, c, q, r, signCode(r), multibitCode(r, 4)); + } + + float[] residualQuery() { + float[] w = new float[q.length]; + for (int i = 0; i < q.length; i++) { + w[i] = q[i] - c[i]; + } + return w; + } + } + + private static double trueDistance(VectorDistanceMetric metric, Fixture f) { + switch (metric) { + case L2: { + double s = 0.0; + for (int i = 0; i < f.q.length; i++) { + double d = (double) f.q[i] - f.x[i]; + s += d * d; + } + return s; // squared, matching the finish + } + case DOT_PRODUCT: + return -dot(f.q, f.x); + case COSINE: + return 1.0 - dot(f.q, f.x) / (norm(f.q) * norm(f.x)); + default: + throw new AssertionError(); + } + } + + private static MetricQueryState state(VectorDistanceMetric metric, float[] query, boolean assumeNormalized) { + return MetricQueryState.create(metric, vector -> vector.clone(), query, assumeNormalized); + } + + /** c1_i = +-0.5 following the strict > 0 convention. */ + private static float[] signCode(float[] r) { + float[] code = new float[r.length]; + for (int i = 0; i < r.length; i++) { + code[i] = r[i] > 0f ? 0.5f : -0.5f; + } + return code; + } + + /** Centered B-bit grid code of the unit residual (per-dim symmetric grid). */ + private static float[] multibitCode(float[] r, int bits) { + int dim = r.length; + float[] code = new float[dim]; + double n = norm(r); + if (n < RaBitQNeutralFactors.EPS_NORM) { + return code; + } + float m = 0f; + float[] v = new float[dim]; + for (int i = 0; i < dim; i++) { + v[i] = (float) (r[i] / n); + m = Math.max(m, Math.abs(v[i])); + } + int levels = (1 << bits) - 1; + float dx = 2f * m / levels; + for (int i = 0; i < dim; i++) { + int lvl = Math.max(0, Math.min(levels, Math.round((v[i] + m) / dx))); + code[i] = -m + dx * lvl; + } + return code; + } + + private static float[] gaussian(Random rnd, int dim, float scale) { + float[] v = new float[dim]; + for (int i = 0; i < dim; i++) { + v[i] = (float) rnd.nextGaussian() * scale; + } + return v; + } + + private static double dot(float[] a, float[] b) { + double s = 0.0; + for (int i = 0; i < a.length; i++) { + s += (double) a[i] * b[i]; + } + return s; + } + + private static double norm(float[] v) { + return Math.sqrt(dot(v, v)); + } + + private static void assertRel(double expected, double actual, String msg) { + double denom = Math.max(1.0, Math.abs(expected)); + assertEquals(expected, actual, REL_TOL * denom, msg); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestMetricQueryStateRotation.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestMetricQueryStateRotation.java new file mode 100644 index 0000000000000..421ea59805a9b --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestMetricQueryStateRotation.java @@ -0,0 +1,93 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * Verifies the rotate-once query-math contract (RFC-109 §3): the padded query is rotated exactly + * once, {@code forRotatedCentroid} forms {@code wRot = qRot - cRot} by subtraction (never rotating + * {@code q - c} per cluster), and exact terms are computed in rotated space with inner-product + * invariance under an orthogonal rotation. + */ +public class TestMetricQueryStateRotation { + + /** Orthogonal rotation: swap adjacent coordinate pairs (a permutation, so norm/IP preserving). */ + private static final MetricQueryState.QueryRotation SWAP_PAIRS = v -> { + float[] out = new float[v.length]; + for (int i = 0; i + 1 < v.length; i += 2) { + out[i] = v[i + 1]; + out[i + 1] = v[i]; + } + if (v.length % 2 == 1) { + out[v.length - 1] = v[v.length - 1]; + } + return out; + }; + + private int rotationCalls; + + @Test + void queryIsRotatedExactlyOnceRegardlessOfClusterCount() { + rotationCalls = 0; + MetricQueryState.QueryRotation counting = v -> { + rotationCalls++; + return SWAP_PAIRS.apply(v); + }; + float[] q = {1f, 2f, 3f, 4f}; + MetricQueryState state = + MetricQueryState.create(VectorDistanceMetric.L2, counting, q, false); + // Constructor rotates the query once; probing many clusters must add no further query rotations. + for (int c = 0; c < 100; c++) { + state.forRotatedCentroid(new float[] {0f, 0f, 0f, 0f}); + } + assertEquals(1, rotationCalls, "query must be rotated exactly once"); + } + + @Test + void rotatedQueryMatchesRotationOfQuery() { + float[] q = {1f, 2f, 3f, 4f}; + MetricQueryState state = + MetricQueryState.create(VectorDistanceMetric.L2, SWAP_PAIRS, q, false); + assertArrayEquals(new float[] {2f, 1f, 4f, 3f}, state.rotatedQuery(), 1e-6f); + } + + @Test + void forRotatedCentroidSubtractsInRotatedSpace() { + float[] q = {1f, 2f, 3f, 4f}; + MetricQueryState state = + MetricQueryState.create(VectorDistanceMetric.L2, SWAP_PAIRS, q, false); + float[] rawCentroid = {0.5f, 0.5f, 0.5f, 0.5f}; + float[] cRot = state.rotateCentroid(rawCentroid); + MetricQueryState.ClusterQuery cq = state.forRotatedCentroid(cRot); + // wRot = qRot - cRot = [2,1,4,3] - [0.5,0.5,0.5,0.5] + assertArrayEquals(new float[] {1.5f, 0.5f, 3.5f, 2.5f}, cq.rotatedQuery, 1e-6f); + assertEquals(1.5f + 0.5f + 3.5f + 2.5f, cq.querySum, 1e-6f); + } + + @Test + void innerProductIsInvariantUnderOrthogonalRotation() { + float[] q = {1f, 2f, 3f, 4f}; + float[] rawCentroid = {0.5f, 0.5f, 0.5f, 0.5f}; + MetricQueryState state = + MetricQueryState.create(VectorDistanceMetric.DOT_PRODUCT, SWAP_PAIRS, q, false); + MetricQueryState.ClusterQuery cq = state.forRotatedCentroid(state.rotateCentroid(rawCentroid)); + // must equal raw = 0.5*(1+2+3+4) = 5.0 for an orthogonal rotation. + assertEquals(5.0, cq.terms.qDotC, 1e-6); + } + + @Test + void rotatedCentroidDimensionMismatchThrows() { + float[] q = {1f, 2f, 3f, 4f}; + MetricQueryState state = + MetricQueryState.create(VectorDistanceMetric.L2, SWAP_PAIRS, q, false); + org.junit.jupiter.api.Assertions.assertThrows(IllegalArgumentException.class, + () -> state.forRotatedCentroid(new float[] {0f, 0f})); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQByteLutScorer.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQByteLutScorer.java new file mode 100644 index 0000000000000..5e760ab432dcf --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQByteLutScorer.java @@ -0,0 +1,156 @@ +/* + * 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.hudi.common.index.vector; + +import org.apache.hudi.avro.model.HoodieVectorIndexPostingBlock; + +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.util.Arrays; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** Scalar-to-byte-LUT identity tests for persisted posting-plane ordering. */ +class TestRaBitQByteLutScorer { + + private static final int DIMENSION = 13; + private static final int CODE_ROW_BYTES = Long.BYTES; + private static final float[] QUERY = { + 0.2f, -0.7f, 1.1f, 0.3f, -0.4f, 0.9f, -1.3f, + 0.6f, 0.8f, -0.2f, 0.5f, -0.1f, 0.4f + }; + + @Test + void lutMatchesScalarForAllSupportedBitWidthsAndPersistedPlanes() { + for (int bits : new int[] {1, 2, 4, 8}) { + RaBitQEncoder encoder = new RaBitQEncoder(DIMENSION, bits, 17L, false); + QuantizedVector encoded = encoder.encodeResidual(vector(), center()); + byte[] signRow = padded(encoded.code); + byte[] extendedPlanes = toMsbFirstPlanes(encoded.extendedCode, bits - 1); + + HoodieVectorIndexPostingBlock block = new PostingBlockBuilder(CODE_ROW_BYTES, bits - 1) + .addRow("first", new byte[CODE_ROW_BYTES], new byte[(bits - 1) * CODE_ROW_BYTES], + 0f, 0f, 0f, 0f, 0f, 0f, "fg", "001", "p", 0) + .addRow("target", signRow, extendedPlanes, + 0f, 0f, 0f, 0f, 0f, 0f, "fg", "001", "p", 1) + .build(); + PostingBlockView view = new PostingBlockView(block); + float querySum = sum(QUERY); + RaBitQByteLutScorer scorer = + RaBitQByteLutScorer.forQuery(QUERY, querySum, DIMENSION, CODE_ROW_BYTES); + + int signOffset = view.signPlaneOffset(1); + double signDot = scorer.planeDot(view.signPlaneBuffer(), signOffset); + float expectedPass1 = RaBitQEncoder.dotPackedBinary(encoded.code, QUERY, DIMENSION) + - 0.5f * querySum; + assertEquals(expectedPass1, scorer.pass1FromDot(signDot), 1e-5f, "pass 1, bits=" + bits); + + float expectedPass2 = RaBitQEncoder.multibitDotTerm( + QUERY, querySum, encoded.code, encoded.extendedCode, DIMENSION, bits); + assertEquals(expectedPass2, + scorer.pass2(signDot, view, view.exPlanesBuffer(), 1, bits - 1, bits), + 5e-5f, + "pass 2, bits=" + bits); + } + } + + @Test + void pass2RejectsMismatchedPostingLayoutAndVectorIndex() { + PostingBlockView oneBitView = new PostingBlockView( + new PostingBlockBuilder(CODE_ROW_BYTES, 0) + .addRow("row", new byte[CODE_ROW_BYTES], new byte[0], + 0f, 0f, 0f, 0f, 0f, 0f, "fg", "001", "p", 0) + .build()); + RaBitQByteLutScorer scorer = + RaBitQByteLutScorer.forQuery(QUERY, sum(QUERY), DIMENSION, CODE_ROW_BYTES); + assertThrows(IllegalArgumentException.class, + () -> scorer.pass2(0, oneBitView, oneBitView.exPlanesBuffer(), 1, 0, 1)); + + PostingBlockView widerView = new PostingBlockView( + new PostingBlockBuilder(2 * CODE_ROW_BYTES, 0) + .addRow("row", new byte[2 * CODE_ROW_BYTES], new byte[0], + 0f, 0f, 0f, 0f, 0f, 0f, "fg", "001", "p", 0) + .build()); + assertThrows(IllegalArgumentException.class, + () -> scorer.pass2(0, widerView, widerView.exPlanesBuffer(), 0, 0, 1)); + } + + @Test + void planeDotSupportsNonzeroBufferOffsetsAndRejectsTruncation() { + byte[] code = {(byte) 0b01010101, (byte) 0b00010101}; + byte[] row = padded(code); + ByteBuffer buffer = ByteBuffer.allocate(3 + CODE_ROW_BYTES); + buffer.position(3); + buffer.put(row); + RaBitQByteLutScorer scorer = + RaBitQByteLutScorer.forQuery(QUERY, sum(QUERY), DIMENSION, CODE_ROW_BYTES); + + assertEquals(RaBitQEncoder.dotPackedBinary(code, QUERY, DIMENSION), + scorer.planeDot(buffer, 3), 1e-6); + assertThrows(IllegalArgumentException.class, () -> scorer.planeDot(buffer, 4)); + } + + private static byte[] toMsbFirstPlanes(byte[] packedLevels, int exBits) { + byte[] planes = new byte[exBits * CODE_ROW_BYTES]; + for (int dimension = 0; dimension < DIMENSION; dimension++) { + int level = unpack(packedLevels, dimension * exBits, exBits); + for (int plane = 0; plane < exBits; plane++) { + int sourceBit = exBits - 1 - plane; + if ((level & (1 << sourceBit)) != 0) { + planes[plane * CODE_ROW_BYTES + (dimension >> 3)] |= (byte) (1 << (dimension & 7)); + } + } + } + return planes; + } + + private static int unpack(byte[] packed, int bitOffset, int bitCount) { + int value = 0; + for (int bit = 0; bit < bitCount; bit++) { + int absoluteBit = bitOffset + bit; + if ((packed[absoluteBit >> 3] & (1 << (absoluteBit & 7))) != 0) { + value |= 1 << bit; + } + } + return value; + } + + private static byte[] padded(byte[] code) { + return Arrays.copyOf(code, CODE_ROW_BYTES); + } + + private static float sum(float[] values) { + float result = 0f; + for (float value : values) { + result += value; + } + return result; + } + + private static float[] vector() { + return new float[] {1f, 3f, -2f, 4f, 0.5f, -1f, 2f, 0.1f, -0.4f, 1.7f, 3.2f, -2.1f, 0.8f}; + } + + private static float[] center() { + return new float[] {0.2f, 2f, -1f, 3f, 0f, -0.2f, 1f, 0f, -0.1f, 1f, 2f, -1f, 0.2f}; + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQEncoder.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQEncoder.java new file mode 100644 index 0000000000000..51dfa793aecd0 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQEncoder.java @@ -0,0 +1,372 @@ +/* + * 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.hudi.common.index.vector; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.util.Random; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Unit tests for {@link RaBitQEncoder}. + */ +class TestRaBitQEncoder { + + private static final int DIM = 64; // small for speed + private RaBitQEncoder encoder; + + @BeforeEach + void setUp() { + encoder = new RaBitQEncoder(DIM, 42L, false); + } + + @Test + void codeBytesSizeCorrect() { + assertEquals((DIM + 7) / 8, encoder.codeBytes()); + } + + @Test + void encodeProducesCorrectCodeLength() { + float[] v = randomVector(DIM, 0); + QuantizedVector encoded = encoder.encode(v); + assertEquals(encoder.codeBytes(), encoded.code.length); + } + + @Test + void encodeScalarIsNorm() { + float[] v = {3f, 4f}; + RaBitQEncoder enc2 = new RaBitQEncoder(2, 42L, false); + QuantizedVector encoded = enc2.encode(v); + assertEquals(5f, encoded.scalar, 1e-4f); + } + + @Test + void encodeAssumeNormalizedScalarIsOne() { + RaBitQEncoder enc = new RaBitQEncoder(DIM, 42L, true); + float[] v = randomVector(DIM, 1); + QuantizedVector encoded = enc.encode(v); + assertEquals(1.0f, encoded.scalar, 1e-6f); + } + + @Test + void deterministic_sameSeedSameCode() { + float[] v = randomVector(DIM, 99); + RaBitQEncoder enc1 = new RaBitQEncoder(DIM, 42L, false); + RaBitQEncoder enc2 = new RaBitQEncoder(DIM, 42L, false); + assertArrayEquals(enc1.encode(v).code, enc2.encode(v).code); + } + + @Test + void differentSeedDifferentCode() { + float[] v = randomVector(DIM, 99); + RaBitQEncoder enc1 = new RaBitQEncoder(DIM, 42L, false); + RaBitQEncoder enc2 = new RaBitQEncoder(DIM, 7L, false); + byte[] c1 = enc1.encode(v).code; + byte[] c2 = enc2.encode(v).code; + // Very unlikely to be equal for a random vector with different rotation seeds + assertFalse(java.util.Arrays.equals(c1, c2), + "Different seeds should produce different codes for a random vector"); + } + + @Test + void hammingDistanceSelf() { + byte[] code = {(byte) 0b10110101, (byte) 0b11001100}; + assertEquals(0, RaBitQDistanceScorer.hammingDistance(code, code)); + } + + @Test + void hammingDistanceKnown() { + byte[] a = {(byte) 0xFF}; + byte[] b = {(byte) 0x00}; + assertEquals(8, RaBitQDistanceScorer.hammingDistance(a, b)); + } + + @Test + void estimateDistanceSameVector() { + float[] v = randomVector(DIM, 5); + QuantizedVector encoded = encoder.encode(v); + RaBitQQueryState qs = encoder.encodeQuery(v); + // Same vector → Hamming distance should be very small → distance close to 0 + float dist = encoder.estimateDistance(qs, encoded); + assertTrue(dist >= 0f, "Distance must be non-negative"); + assertTrue(dist < 0.2f, "Same-vector distance should be very small, got: " + dist); + } + + @Test + void estimateDistanceOppositeVector() { + float[] v = randomVector(DIM, 5); + float[] neg = new float[DIM]; + for (int i = 0; i < DIM; i++) { + neg[i] = -v[i]; + } + + QuantizedVector encoded = encoder.encode(neg); + RaBitQQueryState qs = encoder.encodeQuery(v); + float dist = encoder.estimateDistance(qs, encoded); + // Opposite vectors → distance should be larger than same-vector case + assertTrue(dist > 0.5f, "Opposite vectors should have large distance, got: " + dist); + } + + @Test + void rotationMatrixIsOrthogonal() { + // Each row should have unit norm and rows should be orthogonal + float[][] rotMat = RaBitQEncoder.buildRotationMatrix(8, 42L); + for (int i = 0; i < 8; i++) { + double norm = 0.0; + for (float v : rotMat[i]) { + norm += (double) v * v; + } + assertEquals(1.0, norm, 1e-4, "Row " + i + " should have unit norm"); + } + // Check orthogonality of first two rows + double dot = 0.0; + for (int j = 0; j < 8; j++) { + dot += (double) rotMat[0][j] * rotMat[1][j]; + } + assertEquals(0.0, dot, 1e-4, "First two rows should be orthogonal"); + } + + @Test + void wrongDimensionThrows() { + assertThrows(IllegalArgumentException.class, + () -> encoder.encode(new float[]{1f, 2f})); + } + + @Test + void rejectsNonFiniteVectorsQueriesAndCentroids() { + RaBitQEncoder twoDimensional = new RaBitQEncoder(2); + assertThrows(IllegalArgumentException.class, + () -> twoDimensional.encode(new float[] {Float.NaN, 1f})); + assertThrows(IllegalArgumentException.class, + () -> twoDimensional.encodeQuery(new float[] {1f, Float.POSITIVE_INFINITY})); + assertThrows(IllegalArgumentException.class, + () -> twoDimensional.encodeResidual(new float[] {1f, 2f}, new float[] {1f})); + } + + @Test + void oneDimensionalResidualProducesFiniteExactFactors() { + QuantizedVector encoded = new RaBitQEncoder(1, 4, 42L, false) + .encodeResidual(new float[] {2f}, new float[] {0f}); + assertEquals(0f, encoded.error1, 0f); + assertTrue(Float.isFinite(encoded.rescaleFactor1)); + } + + @Test + void scorerIgnoresPackedPaddingBits() { + assertEquals(1f, RaBitQDistanceScorer.symmetricCosine( + new byte[] {0b00000001}, new byte[] {(byte) 0b11111111}, 1), 0f); + } + + @Test + void quantizedVectorRejectsInvalidPayloadsAndCopiesCodes() { + assertThrows(IllegalArgumentException.class, + () -> new QuantizedVector(new byte[] {1}, null, 1f, null, null, 2)); + assertThrows(IllegalArgumentException.class, + () -> new QuantizedVector(new byte[] {1}, new byte[] {1}, 1f, null, null, 1)); + assertThrows(IllegalArgumentException.class, + () -> new QuantizedVector(new byte[] {1}, new byte[0], 1f, + 0f, 0f, 0f, 0f, -1f, 1f, 1)); + + byte[] code = {1}; + QuantizedVector vector = new QuantizedVector(code, 1f); + code[0] = 2; + assertArrayEquals(new byte[] {1}, vector.getCode()); + byte[] returned = vector.getCode(); + returned[0] = 3; + assertArrayEquals(new byte[] {1}, vector.getCode()); + } + + @Test + void normKnownValue() { + assertEquals(5f, RaBitQEncoder.norm(new float[]{3f, 4f}), 1e-5f); + assertEquals(0f, RaBitQEncoder.norm(new float[]{0f, 0f}), 1e-5f); + } + + // ---- metric-aware reconstruction (RFC-109 metric-aware scoring) ---------------------- + + @Test + void reconstructDistanceL2CollinearRanksByMagnitude() { + // Worked example: query ||q||=10, collinear candidates (cos=1). + // Legacy cosine proxy tied both at 0; correct L2 ranks by magnitude difference. + float qNorm = 10f; + float distSame = RaBitQDistanceScorer.reconstructDistance(VectorDistanceMetric.L2, 1.0f, qNorm, 10f); + float distTiny = RaBitQDistanceScorer.reconstructDistance(VectorDistanceMetric.L2, 1.0f, qNorm, 0.1f); + assertEquals(0f, distSame, 1e-3f, "Identical-magnitude collinear vector should be L2 distance 0"); + assertEquals(9.9f, distTiny, 1e-2f, "Tiny collinear vector should be ~9.9 away"); + assertTrue(distSame < distTiny, "L2 must rank same-magnitude closer than tiny collinear"); + } + + @Test + void reconstructDistanceCosineIgnoresMagnitude() { + assertEquals(0f, RaBitQDistanceScorer.reconstructDistance(VectorDistanceMetric.COSINE, 1.0f, 10f, 10f), 1e-5f); + assertEquals(0f, RaBitQDistanceScorer.reconstructDistance(VectorDistanceMetric.COSINE, 1.0f, 10f, 0.1f), 1e-5f); + assertEquals(2f, RaBitQDistanceScorer.reconstructDistance(VectorDistanceMetric.COSINE, -1.0f, 1f, 1f), 1e-5f); + } + + @Test + void reconstructDistanceDotProductNegatesSimilarity() { + assertEquals(-100f, RaBitQDistanceScorer.reconstructDistance(VectorDistanceMetric.DOT_PRODUCT, 1.0f, 10f, 10f), 1e-3f); + } + + @Test + void metricAwareL2EstimateRanksTrueNeighborFirst() { + // End-to-end: an L2 near neighbour must score below a collinear-but-far vector. + // The legacy cosine scoring scored them identically (the root-cause bug). + int dim = 64; + RaBitQEncoder enc = new RaBitQEncoder(dim, 42L, false); + float[] q = randomVector(dim, 5); + float[] a = q.clone(); // same direction, same magnitude -> true L2 ~0 + float[] b = new float[dim]; // same direction, 10x magnitude -> far in L2, identical cosine + for (int i = 0; i < dim; i++) { + b[i] = q[i] * 10f; + } + RaBitQQueryState qs = enc.encodeQuery(q); + float dA = enc.estimateDistance(qs, enc.encode(a), VectorDistanceMetric.L2, false); + float dB = enc.estimateDistance(qs, enc.encode(b), VectorDistanceMetric.L2, false); + assertTrue(dA < dB, + "L2 must rank same-magnitude neighbour (" + dA + ") closer than 10x-magnitude (" + dB + ")"); + } + + @Test + void asymmetricCosineSameVectorIsHigh() { + int dim = 64; + RaBitQEncoder enc = new RaBitQEncoder(dim, 42L, false); + float[] v = randomVector(dim, 7); + RaBitQQueryState qs = enc.encodeQuery(v); + float cos = RaBitQDistanceScorer.asymmetricCosine(qs.getRotatedQuery(), enc.encode(v).code, dim); + assertTrue(cos > 0.6f, "Asymmetric cosine of a vector with itself should be high, got: " + cos); + } + + @Test + void encodeQueryCapturesQueryNorm() { + RaBitQEncoder enc2 = new RaBitQEncoder(2, 42L, false); + RaBitQQueryState qs = enc2.encodeQuery(new float[]{3f, 4f}); + assertEquals(5f, qs.getQueryNorm(), 1e-4f); + } + + @Test + void encodeResidualFactorsMatchNeutralCodeIdentity() { + int dim = 16; + int bits = 3; + long seed = 101L; + float[] vector = randomVector(dim, 11); + float[] center = randomVector(dim, 13); + RaBitQEncoder enc = new RaBitQEncoder(dim, bits, seed, false); + QuantizedVector encoded = enc.encodeResidual(vector, center); + + float[][] rotation = RaBitQEncoder.buildRotationMatrix(dim, seed); + float[] rotatedVector = rotate(rotation, vector); + float[] rotatedCenter = rotate(rotation, center); + float[] residual = subtract(rotatedVector, rotatedCenter); + double residualSq = dot(residual, residual); + double centerRip = dot(rotatedCenter, residual); + + float[] signCode = signOnlyCode(encoded.code, dim); + double signIpResidual = dot(residual, signCode); + assertRelativeEquals(centerRip, encoded.additiveFactor1, 1.0e-4d); + assertRelativeEquals(residualSq / signIpResidual, encoded.rescaleFactor1, 1.0e-4d); + + float[] fullCode = fullCenteredCode(encoded.code, encoded.extendedCode, dim, bits); + double fullIpResidual = dot(residual, fullCode); + assertRelativeEquals(centerRip, encoded.additiveFactor, 1.0e-4d); + assertRelativeEquals(residualSq / fullIpResidual, encoded.rescaleFactor, 1.0e-4d); + assertRelativeEquals(Math.sqrt(residualSq), encoded.scalar, 1.0e-4d); + assertRelativeEquals(RaBitQEncoder.norm(vector), encoded.vectorNorm, 1.0e-4d); + } + + // ---- helpers ----------------------------------------------------------- + + private static float[] randomVector(int dim, int seed) { + Random rng = new Random(seed); + float[] v = new float[dim]; + for (int i = 0; i < dim; i++) { + v[i] = (float) rng.nextGaussian(); + } + return v; + } + + private static float[] rotate(float[][] rotation, float[] vector) { + float[] out = new float[vector.length]; + for (int i = 0; i < vector.length; i++) { + double acc = 0.0d; + for (int j = 0; j < vector.length; j++) { + acc += (double) rotation[i][j] * vector[j]; + } + out[i] = (float) acc; + } + return out; + } + + private static float[] subtract(float[] left, float[] right) { + float[] out = new float[left.length]; + for (int i = 0; i < left.length; i++) { + out[i] = left[i] - right[i]; + } + return out; + } + + private static double dot(float[] left, float[] right) { + double sum = 0.0d; + for (int i = 0; i < left.length; i++) { + sum += (double) left[i] * right[i]; + } + return sum; + } + + private static float[] signOnlyCode(byte[] signCode, int dimension) { + float[] code = new float[dimension]; + for (int i = 0; i < dimension; i++) { + code[i] = (signCode[i >> 3] & (1 << (i & 7))) != 0 ? 0.5f : -0.5f; + } + return code; + } + + private static float[] fullCenteredCode(byte[] signCode, byte[] extendedCode, int dimension, int bits) { + int exBits = bits - 1; + float cBias = (float) -((1 << exBits) - 0.5d); + float[] code = new float[dimension]; + int bitOffset = 0; + for (int i = 0; i < dimension; i++) { + int level = 0; + for (int bit = 0; bit < exBits; bit++) { + int absoluteBit = bitOffset + bit; + if ((extendedCode[absoluteBit >> 3] & (1 << (absoluteBit & 7))) != 0) { + level |= 1 << bit; + } + } + boolean positive = (signCode[i >> 3] & (1 << (i & 7))) != 0; + code[i] = level + (positive ? (1 << exBits) : 0) + cBias; + bitOffset += exBits; + } + return code; + } + + private static void assertRelativeEquals(double expected, double actual, double tolerance) { + double scale = Math.max(1.0d, Math.abs(expected)); + assertTrue(Math.abs(expected - actual) <= tolerance * scale, + "expected=" + expected + ", actual=" + actual); + } +} \ No newline at end of file diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQNeutralFactors.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQNeutralFactors.java new file mode 100644 index 0000000000000..6dad0bcf2fb6d --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQNeutralFactors.java @@ -0,0 +1,139 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Verifies the corrected RaBitQ neutral-factor semantics (RFC-109 §3): posting-block-owned factor layout, + * removal of the invalid {@code min(1, eps1)} cap, the normalized {@code gMin} quality gate, and + * the exact-zero vs small-nonzero residual split. These are the degenerate cases the previous + * implementation handled incorrectly. + */ +public class TestRaBitQNeutralFactors { + + private static final RaBitQFactorConfig CFG = RaBitQFactorConfig.defaults(); + private static final int D = 4; + + /** code1 for a residual: +0.5 where r>0 else -0.5. */ + private static float[] signCode(float[] r) { + float[] c = new float[r.length]; + for (int i = 0; i < r.length; i++) { + c[i] = r[i] > 0f ? 0.5f : -0.5f; + } + return c; + } + + private static double dot(float[] a, float[] b) { + double s = 0; + for (int i = 0; i < a.length; i++) { + s += (double) a[i] * b[i]; + } + return s; + } + + @Test + void defaultsUseFirstPersistedFactorVersion() { + assertEquals(1, RaBitQFactorConfig.FACTOR_VERSION); + assertEquals(RaBitQFactorConfig.FACTOR_VERSION, CFG.getFactorVersion()); + assertThrows(IllegalArgumentException.class, + () -> new RaBitQFactorConfig(2, 1.9, 1.0e-3, 1.0, 1.0e-3)); + } + + @Test + void exactZeroResidualUsesZeroError() { + float[] r = {0f, 0f, 0f, 0f}; + float[] x = {5f, 0f, 0f, 0f}; + RaBitQNeutralFactors.Factors f = + RaBitQNeutralFactors.compute(r, new float[D], x, 0.0, 0.0, D, CFG); + assertEquals(RaBitQFactorConfig.FACTOR_VERSION, f.factorVersion); + assertEquals(0f, f.residualNorm, 0f); + assertEquals(0f, f.err1, 0f, "only an exact-zero residual may report ERR_1 = 0"); + assertEquals(0f, f.fRescale1, 0f); + } + + @Test + void oneDimensionalResidualUsesExactBound() { + float[] residual = {2f}; + double innerProduct = dot(residual, signCode(residual)); + RaBitQNeutralFactors.Factors factors = RaBitQNeutralFactors.compute( + residual, new float[1], residual, innerProduct, innerProduct, 1, CFG); + + assertEquals(4f, factors.fRescale1, 0f); + assertEquals(0f, factors.err1, 0f); + assertTrue(Float.isFinite(factors.err1)); + } + + @Test + void smallNonzeroResidualDisablesEstimatorWithResidualNormBound() { + // residual norm 1e-4 is <= epsNRel(1e-3) * ||x||(100) = 0.1 -> small-nonzero tier. + float[] r = {1.0e-4f, 0f, 0f, 0f}; + float[] x = {100f, 0f, 0f, 0f}; + double ip1 = dot(r, signCode(r)); + RaBitQNeutralFactors.Factors f = + RaBitQNeutralFactors.compute(r, new float[D], x, ip1, ip1, D, CFG); + assertEquals(0f, f.fRescale1, 0f, "small residual must disable the pass-1 estimator"); + assertEquals(1.0e-4f, f.err1, 1e-9f, "ERR_1 must equal the residual norm, not 0"); + assertTrue(f.err1 > 0f, "a nonzero residual must not report zero error"); + } + + @Test + void largeRelativeErrorDisablesEstimatorInsteadOfClamping() { + // r=[1,0,0,0]: gHat1 = 0.5 -> eps1 = 1.9*sqrt(0.75/0.75) = 1.9 > eps1Max(1.0). + // Old code clamped to min(1,eps1)=1 (invalid); the corrected implementation disables the estimator. + float[] r = {1f, 0f, 0f, 0f}; + float[] x = {1f, 0f, 0f, 0f}; + double ip1 = dot(r, signCode(r)); // = 0.5 + RaBitQNeutralFactors.Factors f = + RaBitQNeutralFactors.compute(r, new float[D], x, ip1, ip1, D, CFG); + assertEquals(0f, f.fRescale1, 0f, "eps1 > eps1Max must disable the estimator"); + assertEquals(1.0f, f.err1, 1e-6f, "ERR_1 must be the maximal valid bound n, not the clamped 1*... "); + } + + @Test + void weakAlignmentBelowGMinIsGated() { + // Drive a near-orthogonal code via an explicitly tiny ipResidual1 (gHat1 << gMin). + float[] r = {1f, 0f, 0f, 0f}; + float[] x = {1f, 0f, 0f, 0f}; + double tinyIp = 1.0e-9; // gHat1 = 1e-9 / (1 * 0.5 * 2) = 1e-9 < gMin(1e-3) + RaBitQNeutralFactors.Factors f = + RaBitQNeutralFactors.compute(r, new float[D], x, tinyIp, tinyIp, D, CFG); + assertEquals(0f, f.fRescale1, 0f, "gHat1 < gMin must disable the estimator"); + assertEquals(1.0f, f.err1, 1e-6f, "gated estimator uses ERR_1 = residualNorm"); + } + + @Test + void perfectAlignmentGivesZeroPass1Error() { + // r = [1,1,1,1]: gHat1 = 1.0 (sign code perfectly aligned) -> eps1 = 0 -> ERR_1 = 0, enabled. + float[] r = {1f, 1f, 1f, 1f}; + float[] x = {1f, 1f, 1f, 1f}; + double ip1 = dot(r, signCode(r)); // = 2.0 + RaBitQNeutralFactors.Factors f = + RaBitQNeutralFactors.compute(r, new float[D], x, ip1, ip1, D, CFG); + assertEquals(2.0f, f.fRescale1, 1e-6f, "F_RESCALE_1 = n^2/ip = 4/2"); + assertEquals(0f, f.err1, 1e-6f, "perfect alignment yields zero pass-1 error"); + } + + @Test + void wellConditionedResidualEnablesEstimatorAndReproducesRip() { + // r = [3,1,1,1]: gHat1 = 0.866 -> eps1 = 0.63 in (0,1); ERR_1 in (0, n). + float[] r = {3f, 1f, 1f, 1f}; + float[] x = {3f, 1f, 1f, 1f}; + float[] code1 = signCode(r); // all +0.5 + double ip1 = dot(r, code1); // = 3.0 + RaBitQNeutralFactors.Factors f = + RaBitQNeutralFactors.compute(r, new float[D], x, ip1, ip1, D, CFG); + assertTrue(f.fRescale1 != 0f, "well-conditioned residual must keep the estimator enabled"); + // F_RESCALE_1 = n^2 / ipResidual1 = 12 / 3 = 4.0 + assertEquals(4.0f, f.fRescale1, 1e-6f); + assertTrue(f.err1 > 0f && f.err1 < f.residualNorm, + "enabled estimator error must be positive and below the maximal bound n"); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQPlaneKernel.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQPlaneKernel.java new file mode 100644 index 0000000000000..3bc0faa74ba32 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQPlaneKernel.java @@ -0,0 +1,112 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +import org.junit.jupiter.api.Test; + +import java.util.Random; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Scalar-vs-popcount identity for the direct plane kernel (RFC-109 §3). Proves the + * {@link RaBitQPlaneKernel} popcount result equals a scalar reference dot of the same quantized + * query {@code qPrime} against the same centered data code, for both pass 1 (sign only) and pass 2 + * (full multibit). + */ +public class TestRaBitQPlaneKernel { + + private static long[] planeFromBits(boolean[] bits, int words) { + long[] p = new long[words]; + for (int i = 0; i < bits.length; i++) { + if (bits[i]) { + p[i >> 6] |= 1L << (i & 63); + } + } + return p; + } + + @Test + void pass2PopcountEqualsScalarReference() { + for (int trial = 0; trial < 50; trial++) { + Random rng = new Random(1000 + trial); + int dim = 100; // deliberately not a multiple of 64 + int bq = 4; + int exBits = 3; // bits = 4 + int words = (dim + 63) >> 6; + + // Random query and its plane quantization. + float[] query = new float[dim]; + for (int i = 0; i < dim; i++) { + query[i] = (float) (rng.nextGaussian() * 3.0); + } + VectorQueryPlanes q = VectorQueryPlanes.quantize(query, bq); + float[] qPrime = q.reconstruct(); + + // Random data code: sign bit + exLevel in [0, 2^exBits - 1]. + boolean[] signBit = new boolean[dim]; + int[] exLevel = new int[dim]; + boolean[][] exBitsArr = new boolean[exBits][dim]; + for (int i = 0; i < dim; i++) { + signBit[i] = rng.nextBoolean(); + exLevel[i] = rng.nextInt(1 << exBits); + for (int b = 0; b < exBits; b++) { + exBitsArr[b][i] = ((exLevel[i] >> b) & 1) != 0; + } + } + long[] signPlane = planeFromBits(signBit, words); + long[][] exPlanes = new long[exBits][]; + for (int b = 0; b < exBits; b++) { + exPlanes[b] = planeFromBits(exBitsArr[b], words); + } + + // Scalar reference: centeredCode[i] = 2^exBits*sign + exLevel + cBias. + double cBias = -((double) ((1 << (exBits + 1)) - 1)) / 2.0; + double refPass2 = 0.0; + double refPass1 = 0.0; + for (int i = 0; i < dim; i++) { + double centered = (1 << exBits) * (signBit[i] ? 1.0 : 0.0) + exLevel[i] + cBias; + refPass2 += qPrime[i] * centered; + refPass1 += qPrime[i] * ((signBit[i] ? 1.0 : 0.0) - 0.5); + } + + double gotPass2 = RaBitQPlaneKernel.scorePass2(q, signPlane, exPlanes, exBits); + double gotPass1 = RaBitQPlaneKernel.scorePass1(q, signPlane); + + double tol = 1e-6 * Math.max(1.0, Math.abs(refPass2)); + assertEquals(refPass2, gotPass2, tol, "pass2 popcount != scalar at trial " + trial); + assertEquals(refPass1, gotPass1, 1e-6 * Math.max(1.0, Math.abs(refPass1)), + "pass1 popcount != scalar at trial " + trial); + } + } + + @Test + void quantizeReconstructRoundTripsWithinOneStep() { + Random rng = new Random(7); + int dim = 64; + float[] query = new float[dim]; + for (int i = 0; i < dim; i++) { + query[i] = (float) (rng.nextGaussian()); + } + VectorQueryPlanes q = VectorQueryPlanes.quantize(query, 4); + float[] qp = q.reconstruct(); + for (int i = 0; i < dim; i++) { + assertTrue(Math.abs(query[i] - qp[i]) <= q.deltaQ() + 1e-6, + "reconstruction error exceeds one quantization step at i=" + i); + } + } + + @Test + void planeErrorBoundCombinesFactorAndQuantizationError() { + float[] query = {1f, 2f, 3f, 4f, 5f, 6f, 7f, 8f}; + VectorQueryPlanes q = VectorQueryPlanes.quantize(query, 4); + double err1 = 0.1; + double residualNorm = 2.0; + double expected = err1 * q.queryPrimeNorm() + residualNorm * q.quantizationErrorNorm(); + assertEquals(expected, q.planeErrorBound(err1, residualNorm), 1e-9); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQResidualHypothesis.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQResidualHypothesis.java new file mode 100644 index 0000000000000..4a3683c7ff180 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQResidualHypothesis.java @@ -0,0 +1,224 @@ +/* + * 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.hudi.common.index.vector; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; + +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * RFC-109 validation experiment (pure JVM, no Spark / no cluster). + * + *

Goal: prove that approximate-only RaBitQ recall on L2 / large-norm data + * (BigANN/SIFT-like) is broken when we quantize the full unit vector (current + * implementation), and is fixed by quantizing the IVF residual + * {@code x - centroid[cluster]} instead. + * + *

The math: {@code L2^2 = ||q||^2 + ||x||^2 - 2*||q||*||x||*cos}. With SIFT-scale + * norms (~1000) a 1-bit cosine-estimate error of +/-0.1 becomes an L2^2 error of + * {@code ~2*1000^2*0.1 = 200k}, which dwarfs the true gaps between near neighbors -> + * random ranking. Residuals are small & zero-centered, so the same 1-bit code resolves + * them, and {@code ||q_res - x_res|| = ||(q-c) - (x-c)|| = ||q - x||} is recovered + * accurately. + * + *

This experiment isolates ENCODING QUALITY: both methods probe ALL clusters + * (no IVF pruning), so the only difference measured is full-vector vs residual. + */ +public class TestRaBitQResidualHypothesis { + + private static final int DIM = 128; // SIFT dimension + private static final int NUM_CLUSTERS = 50; + private static final int PER_CLUSTER = 200; + private static final int N = NUM_CLUSTERS * PER_CLUSTER; // 10k base vectors + private static final int NUM_QUERIES = 100; + private static final int K = 10; // recall@10 + private static final long SEED = 7L; + + @Test + public void residualEncodingFixesL2RecallOnLargeNormData() { + Random rng = new Random(SEED); + + // ---- 1. Synthetic SIFT-like data: clustered, NON-NEGATIVE, LARGE norms ---- + float[][] centers = new float[NUM_CLUSTERS][DIM]; + for (int c = 0; c < NUM_CLUSTERS; c++) { + for (int d = 0; d < DIM; d++) { + centers[c][d] = 20f + rng.nextFloat() * 160f; // [20,180] -> norm ~ 1130 + } + } + float[][] data = new float[N][DIM]; + int[] assign = new int[N]; + int idx = 0; + for (int c = 0; c < NUM_CLUSTERS; c++) { + for (int p = 0; p < PER_CLUSTER; p++) { + float[] v = new float[DIM]; + for (int d = 0; d < DIM; d++) { + v[d] = Math.max(0f, centers[c][d] + (float) (rng.nextGaussian() * 15.0)); // residual ~170 + } + data[idx] = v; + assign[idx] = c; + idx++; + } + } + + // ---- 2. Encoders / encodings ---- + RaBitQEncoder enc = new RaBitQEncoder(DIM, 42L, false); + + // FULL: encode each full vector (current production behavior) + QuantizedVector[] fullCodes = new QuantizedVector[N]; + for (int i = 0; i < N; i++) { + fullCodes[i] = enc.encode(data[i]); + } + + // RESIDUAL: encode x - centroid[assign(x)] + QuantizedVector[] resCodes = new QuantizedVector[N]; + for (int i = 0; i < N; i++) { + resCodes[i] = enc.encode(sub(data[i], centers[assign[i]])); + } + + // ---- 3. Queries: held-out points generated the same way (have true neighbors) ---- + double recallFullSym = 0; + double recallFullAsym = 0; + double recallResSym = 0; + double recallResAsym = 0; + + for (int qi = 0; qi < NUM_QUERIES; qi++) { + int c = rng.nextInt(NUM_CLUSTERS); + float[] q = new float[DIM]; + for (int d = 0; d < DIM; d++) { + q[d] = Math.max(0f, centers[c][d] + (float) (rng.nextGaussian() * 15.0)); + } + + // ground truth: exact L2 top-K over all N + int[] truth = exactTopK(q, data, K); + + // ---- FULL method ---- + RaBitQQueryState qsFull = (RaBitQQueryState) enc.encodeQuery(q); + float[] distFullSym = new float[N]; + float[] distFullAsym = new float[N]; + for (int i = 0; i < N; i++) { + float symCos = RaBitQDistanceScorer.symmetricCosine(qsFull.getBinaryCode(), fullCodes[i].code, DIM); + float asymCos = RaBitQDistanceScorer.asymmetricCosine(qsFull.getRotatedQuery(), fullCodes[i].code, DIM); + distFullSym[i] = RaBitQDistanceScorer.reconstructDistance( + VectorDistanceMetric.L2, symCos, qsFull.getQueryNorm(), fullCodes[i].scalar); + distFullAsym[i] = RaBitQDistanceScorer.reconstructDistance( + VectorDistanceMetric.L2, asymCos, qsFull.getQueryNorm(), fullCodes[i].scalar); + } + recallFullSym += recall(topKByDist(distFullSym, K), truth); + recallFullAsym += recall(topKByDist(distFullAsym, K), truth); + + // ---- RESIDUAL method (probe ALL clusters; query residual is per-cluster) ---- + RaBitQQueryState[] qResByCluster = new RaBitQQueryState[NUM_CLUSTERS]; + for (int cc = 0; cc < NUM_CLUSTERS; cc++) { + qResByCluster[cc] = (RaBitQQueryState) enc.encodeQuery(sub(q, centers[cc])); + } + float[] distResSym = new float[N]; + float[] distResAsym = new float[N]; + for (int i = 0; i < N; i++) { + RaBitQQueryState qsRes = qResByCluster[assign[i]]; + float symCos = RaBitQDistanceScorer.symmetricCosine(qsRes.getBinaryCode(), resCodes[i].code, DIM); + float asymCos = RaBitQDistanceScorer.asymmetricCosine(qsRes.getRotatedQuery(), resCodes[i].code, DIM); + distResSym[i] = RaBitQDistanceScorer.reconstructDistance( + VectorDistanceMetric.L2, symCos, qsRes.getQueryNorm(), resCodes[i].scalar); + distResAsym[i] = RaBitQDistanceScorer.reconstructDistance( + VectorDistanceMetric.L2, asymCos, qsRes.getQueryNorm(), resCodes[i].scalar); + } + recallResSym += recall(topKByDist(distResSym, K), truth); + recallResAsym += recall(topKByDist(distResAsym, K), truth); + } + + recallFullSym /= NUM_QUERIES; + recallFullAsym /= NUM_QUERIES; + recallResSym /= NUM_QUERIES; + recallResAsym /= NUM_QUERIES; + + System.out.println("================ RFC-109 RESIDUAL HYPOTHESIS ================"); + System.out.printf("Data: N=%d D=%d clusters=%d queries=%d (SIFT-like, L2, non-negative)%n", + N, DIM, NUM_CLUSTERS, NUM_QUERIES); + System.out.printf("FULL-vector recall@%d : symmetric=%.3f asymmetric=%.3f%n", K, recallFullSym, recallFullAsym); + System.out.printf("RESIDUAL recall@%d : symmetric=%.3f asymmetric=%.3f%n", K, recallResSym, recallResAsym); + System.out.println("============================================================"); + + // The hypothesis: residual encoding dramatically beats full-vector on L2/large-norm. + assertTrue(recallResSym > recallFullSym + 0.3, + "Residual symmetric recall should massively beat full-vector; got res=" + recallResSym + " full=" + recallFullSym); + assertTrue(recallResAsym >= recallResSym, + "Asymmetric should be >= symmetric for residual; got asym=" + recallResAsym + " sym=" + recallResSym); + } + + // ---- helpers ---- + + private static float[] sub(float[] a, float[] b) { + float[] out = new float[a.length]; + for (int i = 0; i < a.length; i++) { + out[i] = a[i] - b[i]; + } + return out; + } + + private static float l2sq(float[] a, float[] b) { + float s = 0; + for (int i = 0; i < a.length; i++) { + float d = a[i] - b[i]; + s += d * d; + } + return s; + } + + private static int[] exactTopK(float[] q, float[][] data, int k) { + float[] d = new float[data.length]; + for (int i = 0; i < data.length; i++) { + d[i] = l2sq(q, data[i]); + } + return topKByDist(d, k); + } + + /** Indices of the k smallest distances (simple selection; N small). */ + private static int[] topKByDist(float[] dist, int k) { + Integer[] order = new Integer[dist.length]; + for (int i = 0; i < dist.length; i++) { + order[i] = i; + } + Arrays.sort(order, (x, y) -> Float.compare(dist[x], dist[y])); + int[] out = new int[k]; + for (int i = 0; i < k; i++) { + out[i] = order[i]; + } + return out; + } + + private static double recall(int[] got, int[] truth) { + List t = new ArrayList<>(); + for (int x : truth) { + t.add(x); + } + int hit = 0; + for (int g : got) { + if (t.contains(g)) { + hit++; + } + } + return (double) hit / truth.length; + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQResidualRecall.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQResidualRecall.java new file mode 100644 index 0000000000000..d51b831442b26 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestRaBitQResidualRecall.java @@ -0,0 +1,156 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Random; + +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * End-to-end recall validation of the corrected math path (RFC-109 §3 + §3), simulating a + * re-bootstrap: encodes an IVF-residual multibit index with the new {@link RaBitQNeutralFactors} + * (posting-block format 1), scores queries through the rotate-once {@link MetricQueryState} + + * {@link RaBitQEncoder#multibitDotTerm} estimator, and measures recall@10 against brute-force exact + * L2 truth. + * + *

This is deliberately non-circular: the pass/fail comparison is against exact L2, not against + * the estimator itself. It confirms the rotate-once query math and the corrected factors do not + * regress candidate-set recall on SIFT-like clustered data. + */ +public class TestRaBitQResidualRecall { + + private static final int DIM = 128; + private static final int NUM_CLUSTERS = 50; + private static final int PER_CLUSTER = 200; + private static final int N = NUM_CLUSTERS * PER_CLUSTER; // 10k base vectors + private static final int NUM_QUERIES = 100; + private static final int K = 10; + private static final int BITS = 4; + private static final long SEED = 7L; + + @Test + public void residualMultibitRecallMeetsFloorOnCorrectedPath() { + Random rng = new Random(SEED); + + // SIFT-like clustered, non-negative, large-norm data. + float[][] centers = new float[NUM_CLUSTERS][DIM]; + for (int c = 0; c < NUM_CLUSTERS; c++) { + for (int d = 0; d < DIM; d++) { + centers[c][d] = 20f + rng.nextFloat() * 160f; + } + } + float[][] data = new float[N][DIM]; + int[] assign = new int[N]; + int idx = 0; + for (int c = 0; c < NUM_CLUSTERS; c++) { + for (int p = 0; p < PER_CLUSTER; p++) { + float[] v = new float[DIM]; + for (int d = 0; d < DIM; d++) { + v[d] = Math.max(0f, centers[c][d] + (float) (rng.nextGaussian() * 15.0)); + } + data[idx] = v; + assign[idx] = c; + idx++; + } + } + + // Re-bootstrap the residual multibit index with the corrected factors. + RaBitQEncoder enc = new RaBitQEncoder(DIM, BITS, 42L, false); + QuantizedVector[] codes = new QuantizedVector[N]; + for (int i = 0; i < N; i++) { + codes[i] = enc.encodeResidual(data[i], centers[assign[i]]); + } + // Confirm the new factor layout is actually being produced. + assertTrue(codes[0].rescaleFactor != null && codes[0].additiveFactor != null, + "residual multibit encoding must produce neutral factors"); + + double recallSum = 0; + for (int qi = 0; qi < NUM_QUERIES; qi++) { + int c = rng.nextInt(NUM_CLUSTERS); + float[] q = new float[DIM]; + for (int d = 0; d < DIM; d++) { + q[d] = Math.max(0f, centers[c][d] + (float) (rng.nextGaussian() * 15.0)); + } + int[] truth = exactTopK(q, data, K); + + // Rotate-once query state; rotate each centroid once (corrected §2 path). + MetricQueryState state = + MetricQueryState.create(VectorDistanceMetric.L2, enc::rotateVector, q, false); + MetricQueryState.ClusterQuery[] cqByCluster = new MetricQueryState.ClusterQuery[NUM_CLUSTERS]; + for (int cc = 0; cc < NUM_CLUSTERS; cc++) { + cqByCluster[cc] = state.forRotatedCentroid(state.rotateCentroid(centers[cc])); + } + + float[] approx = new float[N]; + for (int i = 0; i < N; i++) { + MetricQueryState.ClusterQuery cq = cqByCluster[assign[i]]; + float dotTerm = RaBitQEncoder.multibitDotTerm( + cq.rotatedQuery, cq.querySum, codes[i].code, codes[i].extendedCode, DIM, BITS); + double rip = (codes[i].rescaleFactor == null ? 0.0 : codes[i].rescaleFactor) * (double) dotTerm; + float centerRip = codes[i].additiveFactor == null ? 0f : codes[i].additiveFactor; + float residualNorm = codes[i].scalar; + float vectorNorm = codes[i].vectorNorm == null ? Float.NaN : codes[i].vectorNorm; + approx[i] = (float) state.rankingDistance(rip, centerRip, residualNorm, vectorNorm, cq); + } + recallSum += recall(topKByDist(approx, K), truth); + } + double recall = recallSum / NUM_QUERIES; + System.out.printf("[RFC-109] residual multibit recall@%d (B=%d) = %.3f%n", K, BITS, recall); + + // Floor for 4-bit residual multibit on this easy synthetic corpus (all clusters probed). + assertTrue(recall >= 0.85, + "corrected-path recall@" + K + " regressed below floor: " + recall); + } + + private static float l2sq(float[] a, float[] b) { + float s = 0; + for (int i = 0; i < a.length; i++) { + float d = a[i] - b[i]; + s += d * d; + } + return s; + } + + private static int[] exactTopK(float[] q, float[][] data, int k) { + float[] d = new float[data.length]; + for (int i = 0; i < data.length; i++) { + d[i] = l2sq(q, data[i]); + } + return topKByDist(d, k); + } + + private static int[] topKByDist(float[] dist, int k) { + Integer[] order = new Integer[dist.length]; + for (int i = 0; i < dist.length; i++) { + order[i] = i; + } + Arrays.sort(order, (x, y) -> Float.compare(dist[x], dist[y])); + int[] out = new int[k]; + for (int i = 0; i < k; i++) { + out[i] = order[i]; + } + return out; + } + + private static double recall(int[] got, int[] truth) { + List t = new ArrayList<>(); + for (int x : truth) { + t.add(x); + } + int hit = 0; + for (int g : got) { + if (t.contains(g)) { + hit++; + } + } + return (double) hit / truth.length; + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorDistanceMetric.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorDistanceMetric.java new file mode 100644 index 0000000000000..7f4b745d6b565 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorDistanceMetric.java @@ -0,0 +1,94 @@ +/* + * 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.hudi.common.index.vector; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * Unit tests for {@link VectorDistanceMetric}. + */ +class TestVectorDistanceMetric { + + @Test + void cosineIdenticalVectors() { + float[] v = {1f, 2f, 3f}; + assertEquals(0f, VectorDistanceMetric.COSINE.compute(v, v), 1e-6f); + } + + @Test + void cosineOrthogonal() { + float[] a = {1f, 0f}; + float[] b = {0f, 1f}; + assertEquals(1f, VectorDistanceMetric.COSINE.compute(a, b), 1e-6f); + } + + @Test + void cosineOpposite() { + float[] a = {1f, 0f}; + float[] b = {-1f, 0f}; + assertEquals(2f, VectorDistanceMetric.COSINE.compute(a, b), 1e-6f); + } + + @Test + void l2ZeroDistance() { + float[] v = {3f, 4f}; + assertEquals(0f, VectorDistanceMetric.L2.compute(v, v), 1e-6f); + } + + @Test + void l2KnownDistance() { + float[] a = {0f, 0f}; + float[] b = {3f, 4f}; + assertEquals(5f, VectorDistanceMetric.L2.compute(a, b), 1e-6f); + } + + @Test + void dotProductSameDirection() { + float[] a = {1f, 0f}; + float[] b = {2f, 0f}; + // dot = 2, negated distance = -2 + assertEquals(-2f, VectorDistanceMetric.DOT_PRODUCT.compute(a, b), 1e-6f); + } + + @Test + void dotProductOrthogonal() { + float[] a = {1f, 0f}; + float[] b = {0f, 1f}; + assertEquals(0f, VectorDistanceMetric.DOT_PRODUCT.compute(a, b), 1e-6f); + } + + @Test + void dimensionMismatchThrows() { + float[] a = {1f, 2f}; + float[] b = {1f}; + assertThrows(IllegalArgumentException.class, + () -> VectorDistanceMetric.L2.compute(a, b)); + } + + @Test + void fromStringCaseInsensitive() { + assertEquals(VectorDistanceMetric.COSINE, VectorDistanceMetric.fromString("cosine")); + assertEquals(VectorDistanceMetric.L2, VectorDistanceMetric.fromString("L2")); + assertEquals(VectorDistanceMetric.DOT_PRODUCT, VectorDistanceMetric.fromString("dot_product")); + } +} \ No newline at end of file diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorIndexArbiter.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorIndexArbiter.java new file mode 100644 index 0000000000000..05932de7ff94b --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorIndexArbiter.java @@ -0,0 +1,95 @@ +/* + * 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.hudi.common.index.vector; + +import org.apache.hudi.common.index.vector.VectorIndexArbiter.Decision; +import org.apache.hudi.common.index.vector.VectorIndexArbiter.ExclusionCounts; +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * Unit tests for {@link VectorIndexArbiter} — the RFC-109 finalist arbiter decision table. + */ +class TestVectorIndexArbiter { + + private static HoodieRecordGlobalLocation loc(String partition, String instant, String fileId) { + return new HoodieRecordGlobalLocation(partition, instant, fileId, 42L); + } + + @Test + void hitAndLocationMatchesServes() { + HoodieRecordGlobalLocation current = loc("2024/01", "t100", "fg-1"); + assertEquals(Decision.SERVE, + VectorIndexArbiter.classify("2024/01", "fg-1", "t100", current)); + } + + @Test + void hitButDifferentInstantIsStale() { + // Same file group + partition, newer base instant (e.g. compaction rewrote the slice). + HoodieRecordGlobalLocation current = loc("2024/01", "t200", "fg-1"); + assertEquals(Decision.STALE, + VectorIndexArbiter.classify("2024/01", "fg-1", "t100", current)); + } + + @Test + void hitButDifferentFileGroupIsStale() { + // The cluster-move case: record was updated and now lives in a different file group. + HoodieRecordGlobalLocation current = loc("2024/01", "t200", "fg-9"); + assertEquals(Decision.STALE, + VectorIndexArbiter.classify("2024/01", "fg-1", "t100", current)); + } + + @Test + void hitButDifferentPartitionIsStale() { + HoodieRecordGlobalLocation current = loc("2024/02", "t100", "fg-1"); + assertEquals(Decision.STALE, + VectorIndexArbiter.classify("2024/01", "fg-1", "t100", current)); + } + + @Test + void rliMissIsDeleted() { + assertEquals(Decision.DELETED, + VectorIndexArbiter.classify("2024/01", "fg-1", "t100", null)); + } + + @Test + void nullPostingInstantAgainstRealLocationIsStale() { + // A locator missing its base instant cannot claim positional trust; it is not a delete. + HoodieRecordGlobalLocation current = loc("2024/01", "t100", "fg-1"); + assertEquals(Decision.STALE, + VectorIndexArbiter.classify("2024/01", "fg-1", null, current)); + } + + @Test + void exclusionCountsSplitStaleAndDeleted() { + ExclusionCounts counts = new ExclusionCounts(); + counts.record(Decision.SERVE); + counts.record(Decision.STALE); + counts.record(Decision.STALE); + counts.record(Decision.DELETED); + assertEquals(2L, counts.stale()); + assertEquals(1L, counts.deleted()); + assertEquals(3L, counts.total()); + assertEquals("arbiterExclusions{stale=2, deleted=1}", counts.toString()); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorIndexOptions.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorIndexOptions.java new file mode 100644 index 0000000000000..6e5a643f7db3a --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorIndexOptions.java @@ -0,0 +1,159 @@ +/* + * 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.hudi.common.index.vector; + +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class TestVectorIndexOptions { + + @Test + void testDefaultsAreCanonicalAndComplete() { + assertEquals( + opts( + VectorIndexOptions.METRIC, "cosine", + VectorIndexOptions.QUANTIZER, "IVF_RABITQ", + VectorIndexOptions.NUM_CLUSTERS, "256", + VectorIndexOptions.MAX_ITER, "20", + VectorIndexOptions.RABITQ_BITS, "4", + VectorIndexOptions.RABITQ_SEED, "42", + VectorIndexOptions.RABITQ_ASSUME_NORMALIZED, "false", + VectorIndexOptions.QUERY_NUM_PROBES, "32", + VectorIndexOptions.QUERY_REFINE_FACTOR, "50", + VectorIndexOptions.QUERY_MODE, "exact_rerank", + VectorIndexOptions.QUERY_STALE_POLICY, "fail"), + VectorIndexOptions.validateAndNormalize(opts())); + } + + @Test + void testValuesAreNormalizedForPersistence() { + Map normalized = VectorIndexOptions.validateAndNormalize(opts( + VectorIndexOptions.METRIC, "DOT-PRODUCT", + VectorIndexOptions.QUANTIZER, "ivf-rabitq", + VectorIndexOptions.RABITQ_ASSUME_NORMALIZED, "TRUE", + VectorIndexOptions.QUERY_MODE, "EXACT-RERANK", + VectorIndexOptions.QUERY_STALE_POLICY, "WARN")); + + assertEquals("dot_product", normalized.get(VectorIndexOptions.METRIC)); + assertEquals("IVF_RABITQ", normalized.get(VectorIndexOptions.QUANTIZER)); + assertEquals("true", normalized.get(VectorIndexOptions.RABITQ_ASSUME_NORMALIZED)); + assertEquals("exact_rerank", normalized.get(VectorIndexOptions.QUERY_MODE)); + assertEquals("warn", normalized.get(VectorIndexOptions.QUERY_STALE_POLICY)); + assertThrows( + UnsupportedOperationException.class, + () -> normalized.put(VectorIndexOptions.METRIC, "l2")); + } + + @Test + void testEveryMetricQueryModeAndStalePolicyIsAccepted() { + assertCanonical(VectorIndexOptions.METRIC, "cosine", "cosine"); + assertCanonical(VectorIndexOptions.METRIC, "l2", "l2"); + assertCanonical(VectorIndexOptions.METRIC, "dot_product", "dot_product"); + assertCanonical(VectorIndexOptions.QUERY_MODE, "approximate", "approximate"); + assertCanonical(VectorIndexOptions.QUERY_MODE, "exact_rerank", "exact_rerank"); + assertCanonical(VectorIndexOptions.QUERY_STALE_POLICY, "fail", "fail"); + assertCanonical(VectorIndexOptions.QUERY_STALE_POLICY, "warn", "warn"); + assertCanonical(VectorIndexOptions.QUERY_STALE_POLICY, "fallback", "fallback"); + } + + @Test + void testUnknownRetiredAndMisspelledOptionsAreRejected() { + assertInvalidOption("vector.dimension", "128"); + assertInvalidOption("vector.query.nprobe", "8"); + assertInvalidOption("vector.unknown", "value"); + } + + @Test + void testUnsupportedEnumValuesAreRejectedWithOptionContext() { + assertInvalidValueContainsKey(VectorIndexOptions.METRIC, "manhattan"); + assertInvalidValueContainsKey(VectorIndexOptions.QUANTIZER, "pq"); + assertInvalidValueContainsKey(VectorIndexOptions.QUERY_MODE, "fast-ish"); + assertInvalidValueContainsKey(VectorIndexOptions.QUERY_STALE_POLICY, "ignore"); + assertInvalidValueContainsKey(VectorIndexOptions.RABITQ_ASSUME_NORMALIZED, "yes"); + } + + @Test + void testNumericOptionsAreValidatedWithOptionContext() { + assertCanonical(VectorIndexOptions.RABITQ_BITS, "1", "1"); + assertCanonical(VectorIndexOptions.RABITQ_BITS, "8", "8"); + assertInvalidValueContainsKey(VectorIndexOptions.NUM_CLUSTERS, "0"); + assertInvalidValueContainsKey(VectorIndexOptions.MAX_ITER, "-1"); + assertInvalidValueContainsKey(VectorIndexOptions.QUERY_NUM_PROBES, "0"); + assertInvalidValueContainsKey(VectorIndexOptions.QUERY_REFINE_FACTOR, "-1"); + assertInvalidValueContainsKey(VectorIndexOptions.RABITQ_BITS, "0"); + assertInvalidValueContainsKey(VectorIndexOptions.RABITQ_BITS, "9"); + assertInvalidValueContainsKey(VectorIndexOptions.RABITQ_SEED, "many"); + } + + @Test + void testNumProbesMustNotExceedNumClusters() { + assertCanonical( + opts( + VectorIndexOptions.NUM_CLUSTERS, "32", + VectorIndexOptions.QUERY_NUM_PROBES, "32"), + VectorIndexOptions.QUERY_NUM_PROBES, + "32"); + + IllegalArgumentException error = assertThrows( + IllegalArgumentException.class, + () -> VectorIndexOptions.validateAndNormalize(opts( + VectorIndexOptions.NUM_CLUSTERS, "4", + VectorIndexOptions.QUERY_NUM_PROBES, "5"))); + assertTrue(error.getMessage().contains(VectorIndexOptions.QUERY_NUM_PROBES)); + assertTrue(error.getMessage().contains(VectorIndexOptions.NUM_CLUSTERS)); + } + + private static void assertCanonical(String key, String input, String expected) { + assertCanonical(opts(key, input), key, expected); + } + + private static void assertCanonical( + Map options, String key, String expected) { + assertEquals(expected, VectorIndexOptions.validateAndNormalize(options).get(key)); + } + + private static void assertInvalidOption(String key, String value) { + IllegalArgumentException error = assertThrows( + IllegalArgumentException.class, + () -> VectorIndexOptions.validateAndNormalize(opts(key, value))); + assertTrue(error.getMessage().contains(key)); + } + + private static void assertInvalidValueContainsKey(String key, String value) { + IllegalArgumentException error = assertThrows( + IllegalArgumentException.class, + () -> VectorIndexOptions.validateAndNormalize(opts(key, value))); + assertTrue(error.getMessage().contains(key)); + } + + private static Map opts(String... keyValues) { + Map options = new HashMap<>(); + for (int i = 0; i < keyValues.length; i += 2) { + options.put(keyValues[i], keyValues[i + 1]); + } + return options; + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorIndexPruner.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorIndexPruner.java new file mode 100644 index 0000000000000..76c297957dbfc --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/TestVectorIndexPruner.java @@ -0,0 +1,137 @@ +/* + * 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.hudi.common.index.vector; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Unit tests for {@link VectorIndexPruner}. + */ +class TestVectorIndexPruner { + + /** 4 centroids at the corners of a 2D unit square. */ + private static final float[][] CENTROIDS = { + {1f, 1f}, // cluster 0 + {1f, -1f}, // cluster 1 + {-1f, 1f}, // cluster 2 + {-1f, -1f} // cluster 3 + }; + + private VectorIndexPruner pruner(Map> clusterMap) { + return new VectorIndexPruner(CENTROIDS, clusterMap, VectorDistanceMetric.L2); + } + + @Test + void probeTopOneReturnsCorrectFileGroup() { + Map> map = new HashMap<>(); + map.put(0, new HashSet<>(Collections.singletonList("fg-A"))); + map.put(1, new HashSet<>(Collections.singletonList("fg-B"))); + map.put(2, new HashSet<>(Collections.singletonList("fg-C"))); + map.put(3, new HashSet<>(Collections.singletonList("fg-D"))); + + VectorIndexPruner p = pruner(map); + // Query near cluster 0 (1,1) + Set fgs = p.probe(new float[]{0.9f, 0.9f}, 1); + assertEquals(Collections.singleton("fg-A"), fgs); + } + + @Test + void probeTwoReturnsUnionOfFileGroups() { + Map> map = new HashMap<>(); + map.put(0, new HashSet<>(Arrays.asList("fg-A", "fg-E"))); + map.put(1, new HashSet<>(Collections.singletonList("fg-B"))); + map.put(2, new HashSet<>(Collections.singletonList("fg-C"))); + map.put(3, new HashSet<>(Collections.singletonList("fg-D"))); + + VectorIndexPruner p = pruner(map); + // Query near cluster 1 (1,-1), second nearest is cluster 0 (1,1) + Set fgs = p.probe(new float[]{0.9f, -0.9f}, 2); + assertTrue(fgs.contains("fg-A")); + assertTrue(fgs.contains("fg-B")); + assertTrue(fgs.contains("fg-E")); + assertEquals(3, fgs.size()); + } + + @Test + void probeWithNoMatchingClustersInMapReturnsEmpty() { + Map> map = new HashMap<>(); + // cluster 0 has no file groups registered + VectorIndexPruner p = pruner(map); + Set fgs = p.probe(new float[]{0.9f, 0.9f}, 1); + assertTrue(fgs.isEmpty()); + } + + @Test + void probeCapsProbesToAvailableClusters() { + Map> map = new HashMap<>(); + for (int i = 0; i < 4; i++) { + map.put(i, new HashSet<>(Collections.singletonList("fg-" + i))); + } + VectorIndexPruner p = pruner(map); + // Request more probes than clusters + Set fgs = p.probe(new float[]{0f, 0f}, 100); + assertEquals(4, fgs.size()); + } + + @Test + void emptyCentroidsReturnsEmpty() { + VectorIndexPruner p = new VectorIndexPruner( + new float[0][0], Collections.emptyMap(), VectorDistanceMetric.L2); + assertTrue(p.probe(new float[]{1f}, 1).isEmpty()); + } + + @Test + void buildClusterMapFiltersPartitions() { + List assignments = Arrays.asList( + new Object[]{0, "fg-A", "p=2024"}, + new Object[]{1, "fg-B", "p=2023"}, + new Object[]{0, "fg-C", "p=2023"} + ); + Set filter = Collections.singleton("p=2024"); + Map> map = + VectorIndexPruner.buildClusterMap(assignments, filter); + assertEquals(1, map.size()); + assertTrue(map.get(0).contains("fg-A")); + assertFalse(map.containsKey(1)); + } + + @Test + void buildClusterMapNullFilterIncludesAll() { + List assignments = Arrays.asList( + new Object[]{0, "fg-A", "p=2024"}, + new Object[]{1, "fg-B", "p=2023"} + ); + Map> map = + VectorIndexPruner.buildClusterMap(assignments, null); + assertEquals(2, map.size()); + } +} \ No newline at end of file diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestCommonVectorSearchContinuation.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestCommonVectorSearchContinuation.java new file mode 100644 index 0000000000000..dc4b9a30d9ada --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestCommonVectorSearchContinuation.java @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector.search; + +import org.apache.hudi.common.index.vector.VectorDistanceMetric; +import org.apache.hudi.common.index.vector.VectorStalePolicy; +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class TestCommonVectorSearchContinuation { + + @Test + void drawsAnotherRetainedWindowWithoutRescanning() { + VectorSearchBudget budget = new VectorSearchBudget( + 5000, 2, 2, 4, 10, 1, VectorExecutionMode.LOCAL, 10, DeadlinePolicy.FAIL); + VectorSearchRequest request = new VectorSearchRequest( + "embedding", new float[] {1f}, VectorDistanceMetric.L2, 2, 1, 1, true, + VectorStalePolicy.FALLBACK, "001", budget); + List candidates = Arrays.asList( + candidate("deleted-1", 1), candidate("deleted-2", 2), + candidate("live-1", 3), candidate("live-2", 4)); + + AtomicInteger scans = new AtomicInteger(); + AtomicReference retainedPool = new AtomicReference<>(); + VectorCandidateSource source = (plan, context) -> { + scans.incrementAndGet(); + ListVectorCandidatePool pool = new ListVectorCandidatePool(candidates, budget); + retainedPool.set(pool); + return pool; + }; + RecordIndexLookup lookup = (keys, instant) -> { + java.util.Map locations = new java.util.HashMap<>(); + for (String key : keys) { + if (key.startsWith("live")) { + locations.put(key, new HoodieRecordGlobalLocation("p", "001", "f")); + } + } + return locations; + }; + VectorExactReranker reranker = (tasks, req, snapshot, context) -> { + java.util.ArrayList results = new java.util.ArrayList<>(); + for (VectorFetchTask task : tasks.collectAsList()) { + for (VectorRowRequest row : task.getRequests()) { + results.add(new VectorSearchResult( + row.getRecordKey(), row.getApproximateDistance(), + new HoodieRecordGlobalLocation("p", "001", "f"))); + } + } + return org.apache.hudi.common.data.HoodieListData.eager(results); + }; + + CommonVectorSearchExecutor executor = new CommonVectorSearchExecutor( + ignored -> new VectorSearchSnapshot( + "001", new VectorIndexSnapshot(1, 1, 1, "rot-v1", "quant-v1")), + ignored -> new VectorExecutionDecision( + VectorExecutionMode.LOCAL, VectorExecutionMode.LOCAL, 4, 10, "test"), + source, + new RecordIndexVectorCandidateArbiter(lookup), + new DefaultVectorFetchPlanner(), + reranker); + + List results = executor.execute(request, null).collectAsList(); + + assertEquals(Arrays.asList("live-1", "live-2"), + Arrays.asList(results.get(0).getRecordKey(), results.get(1).getRecordKey())); + assertEquals(1, scans.get(), "continuation must not rescan MDT"); + assertEquals(4, retainedPool.get().consumed(), "second window must come from the retained pool"); + } + + private static VectorCandidate candidate(String key, double distance) { + return new VectorCandidate(key, 1, 0, distance, + new VectorPostingLocator(1, 1, 0, 0, 0, "p", "f", "001", 1)); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestCommonVectorSearchExecutor.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestCommonVectorSearchExecutor.java new file mode 100644 index 0000000000000..bfa75f6e83b3e --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestCommonVectorSearchExecutor.java @@ -0,0 +1,112 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieListData; +import org.apache.hudi.common.index.vector.VectorDistanceMetric; +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * End-to-end wiring test for {@link CommonVectorSearchExecutor} (RFC-109 §11): drives the full + * stage pipeline with a fake candidate source and reranker but the real + * {@link RecordIndexVectorCandidateArbiter} and {@link DefaultVectorFetchPlanner}, asserting the + * snapshot is resolved once, DELETED candidates are dropped end-to-end, and results flow through. + */ +public class TestCommonVectorSearchExecutor { + + private static VectorCandidate candidate(String key, String fileId) { + VectorPostingLocator loc = new VectorPostingLocator(1, 0, 0, 0L, 0, "p", fileId, "001", 7L); + return new VectorCandidate(key, 0, 0, 1.0, loc); + } + + private static VectorSearchRequest request() { + VectorSearchBudget budget = VectorSearchBudget.defaults(3, 5000L); + return new VectorSearchRequest("embedding", new float[] {1f, 2f}, VectorDistanceMetric.L2, + 3, 32, 50, true, null, budget); + } + + @Test + void drivesFullPipelineAndDropsDeletedEndToEnd() { + List scanned = new ArrayList<>(); + scanned.add(candidate("k1", "fileA")); + scanned.add(candidate("k2", "fileA")); + scanned.add(candidate("k3", "fileB")); // will be DELETED via RLI miss + + AtomicBoolean snapshotResolved = new AtomicBoolean(false); + VectorSnapshotResolver resolver = req -> { + snapshotResolved.set(true); + return new VectorSearchSnapshot("001", + new VectorIndexSnapshot(1, 1, 1, "rot-v1", "quant-v1")); + }; + + // Real arbiter with a fake RLI: k1/k2 live & matching (SERVE), k3 absent (DELETED). + Map rli = new HashMap<>(); + rli.put("k1", new HoodieRecordGlobalLocation("p", "001", "fileA")); + rli.put("k2", new HoodieRecordGlobalLocation("p", "001", "fileA")); + RecordIndexLookup lookup = (keys, tableInstant) -> { + Map out = new HashMap<>(); + for (String k : keys) { + if (rli.containsKey(k)) { + out.put(k, rli.get(k)); + } + } + return out; + }; + + VectorCandidateSource source = (plan, ec) -> new ListVectorCandidatePool(scanned, plan.getRequest().getBudget()); + VectorCandidateArbiter arbiter = new RecordIndexVectorCandidateArbiter(lookup); + VectorFetchPlanner planner = new DefaultVectorFetchPlanner(); + // Fake reranker: emit one result per fetched row (distance = approx), preserving live location. + VectorExactReranker reranker = (tasks, req, snap, ec) -> { + List results = new ArrayList<>(); + for (VectorFetchTask task : tasks.collectAsList()) { + for (VectorRowRequest r : task.getRequests()) { + results.add(new VectorSearchResult(r.getRecordKey(), r.getApproximateDistance(), + new HoodieRecordGlobalLocation("p", task.getBaseInstant(), task.getFileId()))); + } + } + return HoodieListData.eager(results); + }; + + CommonVectorSearchExecutor executor = new CommonVectorSearchExecutor( + resolver, new ThresholdVectorExecutionModeSelector(), source, arbiter, planner, reranker); + + List results = executor.execute(request(), null).collectAsList(); + + assertTrue(snapshotResolved.get(), "snapshot must be resolved once at the top of the pipeline"); + assertEquals(2, results.size(), "k3 (DELETED via RLI miss) must be dropped end-to-end"); + List keys = new ArrayList<>(); + for (VectorSearchResult r : results) { + keys.add(r.getRecordKey()); + } + assertTrue(keys.contains("k1") && keys.contains("k2")); + assertTrue(!keys.contains("k3")); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestDefaultExactVectorScorer.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestDefaultExactVectorScorer.java new file mode 100644 index 0000000000000..26af4170095b2 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestDefaultExactVectorScorer.java @@ -0,0 +1,82 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.index.vector.VectorDistanceMetric; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Verifies the default exact scorer (RFC-109 §10): squared-L2 for L2 (order-preserving vs the + * true metric), negated dot for DOT_PRODUCT, and 1-cos for COSINE, all in float64. + */ +public class TestDefaultExactVectorScorer { + + private final ExactVectorScorer scorer = new DefaultExactVectorScorer(); + + @Test + void l2ReturnsSquaredDistance() { + float[] q = {1f, 2f, 3f}; + float[] c = {4f, 6f, 3f}; + // diffs: 3,4,0 -> squared sum = 9 + 16 + 0 = 25 (NOT sqrt=5) + assertEquals(25.0, scorer.distance(q, c, VectorDistanceMetric.L2), 1e-9); + } + + @Test + void l2SquaredPreservesOrderVsTrueMetric() { + float[] q = {0f, 0f}; + float[] near = {1f, 0f}; // true L2 = 1, squared = 1 + float[] far = {3f, 0f}; // true L2 = 3, squared = 9 + double dNear = scorer.distance(q, near, VectorDistanceMetric.L2); + double dFar = scorer.distance(q, far, VectorDistanceMetric.L2); + assertTrue(dNear < dFar, "squared L2 must preserve nearest-neighbor ordering"); + } + + @Test + void dotProductIsNegated() { + float[] q = {1f, 2f, 3f}; + float[] c = {1f, 1f, 1f}; + // dot = 6 -> negated = -6 + assertEquals(-6.0, scorer.distance(q, c, VectorDistanceMetric.DOT_PRODUCT), 1e-9); + } + + @Test + void cosineOfParallelVectorsIsZero() { + float[] q = {1f, 2f, 3f}; + float[] c = {2f, 4f, 6f}; // same direction + assertEquals(0.0, scorer.distance(q, c, VectorDistanceMetric.COSINE), 1e-9); + } + + @Test + void cosineHandlesZeroVector() { + float[] q = {0f, 0f, 0f}; + float[] c = {1f, 2f, 3f}; + assertEquals(1.0, scorer.distance(q, c, VectorDistanceMetric.COSINE), 1e-9); + } + + @Test + void dimensionMismatchThrows() { + assertThrows(IllegalArgumentException.class, + () -> scorer.distance(new float[] {1f, 2f}, new float[] {1f}, VectorDistanceMetric.L2)); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestDefaultVectorExactReranker.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestDefaultVectorExactReranker.java new file mode 100644 index 0000000000000..6b3d2e9e214d9 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestDefaultVectorExactReranker.java @@ -0,0 +1,96 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.data.HoodieListData; +import org.apache.hudi.common.index.vector.VectorDistanceMetric; +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * Verifies {@link DefaultVectorExactReranker} (RFC-109 §10): reads via the injected handle, + * scores with {@link DefaultExactVectorScorer}, and returns the global exact top-K in ascending + * distance order, using an in-memory fake read handle. + */ +public class TestDefaultVectorExactReranker { + + /** In-memory handle: returns a VectorRecord per row request, vector looked up by record key. */ + private static HoodieVectorBatchReadHandleSupplier fakeSupplier(Map vectors) { + return () -> new HoodieVectorBatchReadHandle() { + @Override + public Iterator read(VectorFetchTask task, String recordKeyField, String vectorColumn) { + List out = new ArrayList<>(); + for (VectorRowRequest r : task.getRequests()) { + float[] v = vectors.get(r.getRecordKey()); + if (v != null) { + out.add(new VectorRecord(r.getRecordKey(), v, + new HoodieRecordGlobalLocation(task.getPartitionPath(), task.getBaseInstant(), task.getFileId()))); + } + } + return out.iterator(); + } + + @Override + public void close() { + } + }; + } + + @Test + void returnsGlobalExactTopKInAscendingDistance() { + Map vectors = new HashMap<>(); + vectors.put("k1", new float[] {1f, 0f, 0f}); // squared L2 vs origin = 1 + vectors.put("k2", new float[] {2f, 0f, 0f}); // 4 + vectors.put("k3", new float[] {3f, 0f, 0f}); // 9 + vectors.put("k4", new float[] {0.5f, 0f, 0f}); // 0.25 + + List rows = new ArrayList<>(); + rows.add(new VectorRowRequest("k1", 0L, VectorCandidateState.SERVE, 0.0)); + rows.add(new VectorRowRequest("k2", 1L, VectorCandidateState.SERVE, 0.0)); + rows.add(new VectorRowRequest("k3", 2L, VectorCandidateState.SERVE, 0.0)); + rows.add(new VectorRowRequest("k4", 3L, VectorCandidateState.SERVE, 0.0)); + VectorFetchTask task = new VectorFetchTask("p", "fileA", "/tmp/fileA.parquet", "001", rows); + + VectorSearchBudget budget = VectorSearchBudget.defaults(2, 5000L); + VectorSearchRequest request = new VectorSearchRequest( + "embedding", new float[] {0f, 0f, 0f}, VectorDistanceMetric.L2, 2, 32, 50, true, null, budget); + + DefaultVectorExactReranker reranker = new DefaultVectorExactReranker( + fakeSupplier(vectors), new DefaultExactVectorScorer(), "record_key", "embedding"); + + HoodieData tasks = HoodieListData.eager(java.util.Collections.singletonList(task)); + List results = reranker.rerank(tasks, request, null, null).collectAsList(); + + assertEquals(2, results.size(), "topK=2"); + assertEquals("k4", results.get(0).getRecordKey(), "nearest is k4 (0.25)"); + assertEquals(0.25, results.get(0).getDistance(), 1e-6); + assertEquals("k1", results.get(1).getRecordKey(), "second nearest is k1 (1.0)"); + assertEquals(1.0, results.get(1).getDistance(), 1e-6); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestDefaultVectorFetchPlanner.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestDefaultVectorFetchPlanner.java new file mode 100644 index 0000000000000..3e6218b92f04a --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestDefaultVectorFetchPlanner.java @@ -0,0 +1,106 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.data.HoodieListData; +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Verifies {@link DefaultVectorFetchPlanner} (RFC-109 §8): grouping by live file, positional row + * preservation for SERVE, key-fallback ({@code rowPosition = -1}) for STALE, and exclusion of DELETED. + */ +public class TestDefaultVectorFetchPlanner { + + private static VectorCandidate candidate(String key, int cluster, long rowPos, String partition, String fileId) { + VectorPostingLocator loc = new VectorPostingLocator( + 1, cluster, 0, 0L, 0, partition, fileId, "001", rowPos); + return new VectorCandidate(key, cluster, 0, 1.0, loc); + } + + private static ArbitratedVectorCandidate arb(VectorCandidate c, VectorCandidateState state, + String partition, String fileId) { + HoodieRecordGlobalLocation live = state == VectorCandidateState.DELETED + ? null : new HoodieRecordGlobalLocation(partition, "001", fileId); + return new ArbitratedVectorCandidate(c, state, live); + } + + @Test + void groupsByFileAndPreservesPositionsAndFallback() { + List input = new ArrayList<>(); + // File A: two SERVE (positional) + one STALE (key fallback). + input.add(arb(candidate("k1", 0, 10L, "p", "fileA"), VectorCandidateState.SERVE, "p", "fileA")); + input.add(arb(candidate("k2", 0, 20L, "p", "fileA"), VectorCandidateState.SERVE, "p", "fileA")); + input.add(arb(candidate("k3", 0, 30L, "p", "fileA"), VectorCandidateState.STALE, "p", "fileA")); + // File B: one SERVE. + input.add(arb(candidate("k4", 1, 5L, "p", "fileB"), VectorCandidateState.SERVE, "p", "fileB")); + // DELETED: must be dropped. + input.add(arb(candidate("k5", 1, 7L, "p", "fileB"), VectorCandidateState.DELETED, "p", "fileB")); + + HoodieData data = HoodieListData.eager(input); + List tasks = new DefaultVectorFetchPlanner().plan(data, null, null).collectAsList(); + + Map byFile = new HashMap<>(); + for (VectorFetchTask t : tasks) { + byFile.put(t.getFileId(), t); + } + assertEquals(2, tasks.size(), "expected one task per live file (A, B)"); + assertNull(byFile.get("fileB").getBaseFilePath(), "baseFilePath resolved later by read handle"); + + VectorFetchTask a = byFile.get("fileA"); + assertEquals(3, a.size(), "fileA must contain 3 rows (2 SERVE + 1 STALE), DELETED excluded"); + int positional = 0; + int fallback = 0; + for (VectorRowRequest r : a.getRequests()) { + if (r.getState() == VectorCandidateState.SERVE) { + assertTrue(r.isPositional() && r.getRowPosition() >= 0, "SERVE must keep its row position"); + positional++; + } else if (r.getState() == VectorCandidateState.STALE) { + assertEquals(-1L, r.getRowPosition(), "STALE must drop the row position for key fallback"); + assertTrue(!r.isPositional()); + fallback++; + } + } + assertEquals(2, positional); + assertEquals(1, fallback); + + // DELETED k5 excluded -> fileB has only 1 row. + assertEquals(1, byFile.get("fileB").size()); + } + + @Test + void allDeletedProducesNoTasks() { + List input = new ArrayList<>(); + input.add(arb(candidate("k1", 0, 1L, "p", "fileA"), VectorCandidateState.DELETED, "p", "fileA")); + HoodieData data = HoodieListData.eager(input); + List tasks = new DefaultVectorFetchPlanner().plan(data, null, null).collectAsList(); + assertTrue(tasks.isEmpty(), "all-DELETED input must produce no fetch tasks"); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestRecordIndexVectorCandidateArbiter.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestRecordIndexVectorCandidateArbiter.java new file mode 100644 index 0000000000000..c2ad5f360a661 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestRecordIndexVectorCandidateArbiter.java @@ -0,0 +1,157 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.data.HoodieListData; +import org.apache.hudi.common.index.vector.VectorDistanceMetric; +import org.apache.hudi.common.index.vector.VectorStalePolicy; +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; +import org.apache.hudi.exception.HoodieException; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; + +/** + * Verifies {@link RecordIndexVectorCandidateArbiter} (RFC-109 §7): SERVE when the current RLI + * location matches the posting locator, STALE (preserved) when it differs, DELETED (dropped) on an + * RLI miss, using a fake snapshot-pinned {@link RecordIndexLookup}. + */ +public class TestRecordIndexVectorCandidateArbiter { + + private static VectorCandidate candidate(String key, String partition, String fileId, String instant) { + VectorPostingLocator loc = new VectorPostingLocator( + 1, 0, 0, 0L, 0, partition, fileId, instant, 42L); + return new VectorCandidate(key, 0, 0, 1.0, loc); + } + + private static VectorSearchRequest request(VectorStalePolicy stalePolicy) { + return new VectorSearchRequest("embedding", new float[] {1f}, VectorDistanceMetric.L2, + 2, 1, 1, true, stalePolicy, "001", VectorSearchBudget.defaults(2, 1000)); + } + + private static VectorSearchSnapshot snapshot() { + return new VectorSearchSnapshot("001", new VectorIndexSnapshot(1, 1, 1, "rot-v1", "quant-v1")); + } + + @Test + void classifiesServeStaleDeletedAndDropsDeleted() { + List input = new ArrayList<>(); + input.add(candidate("serveKey", "p", "fileA", "001")); // matches -> SERVE + input.add(candidate("staleKey", "p", "fileA", "001")); // current differs -> STALE + input.add(candidate("deletedKey", "p", "fileA", "001")); // RLI miss -> DELETED (dropped) + + // Snapshot-pinned RLI state: serveKey matches posting, staleKey moved to fileB@002, deletedKey absent. + Map rli = new HashMap<>(); + rli.put("serveKey", new HoodieRecordGlobalLocation("p", "001", "fileA")); + rli.put("staleKey", new HoodieRecordGlobalLocation("p", "002", "fileB")); + RecordIndexLookup lookup = (keys, tableInstant) -> { + Map out = new HashMap<>(); + for (String k : keys) { + if (rli.containsKey(k)) { + out.put(k, rli.get(k)); + } + } + return out; + }; + + HoodieData data = HoodieListData.eager(input); + List result = + new RecordIndexVectorCandidateArbiter(lookup) + .arbitrate(data, request(VectorStalePolicy.FALLBACK), snapshot(), null).collectAsList(); + + Map byKey = new HashMap<>(); + for (ArbitratedVectorCandidate a : result) { + byKey.put(a.getCandidate().getRecordKey(), a); + } + + assertEquals(2, result.size(), "DELETED finalist must be dropped"); + + ArbitratedVectorCandidate serve = byKey.get("serveKey"); + assertEquals(VectorCandidateState.SERVE, serve.getState()); + assertNotNull(serve.getLiveLocation()); + assertEquals("fileA", serve.getLiveLocation().getFileId()); + + ArbitratedVectorCandidate stale = byKey.get("staleKey"); + assertEquals(VectorCandidateState.STALE, stale.getState()); + assertNotNull(stale.getLiveLocation(), "STALE must retain the live location for key fallback"); + assertEquals("fileB", stale.getLiveLocation().getFileId()); + + assertNull(byKey.get("deletedKey"), "DELETED must not appear in output"); + } + + @Test + void failPolicyRejectsStaleCandidate() { + HoodieData data = HoodieListData.eager( + java.util.Collections.singletonList(candidate("stale", "p", "old", "001"))); + RecordIndexLookup lookup = (keys, tableInstant) -> java.util.Collections.singletonMap( + "stale", new HoodieRecordGlobalLocation("p", "002", "new")); + + assertThrows(HoodieException.class, () -> new RecordIndexVectorCandidateArbiter(lookup) + .arbitrate(data, request(VectorStalePolicy.FAIL), snapshot(), null).collectAsList()); + } + + @Test + void batchesLookupsAtPinnedInstant() { + List input = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + input.add(candidate("k" + i, "p", "f", "001")); + } + AtomicInteger calls = new AtomicInteger(); + RecordIndexLookup lookup = (keys, tableInstant) -> { + assertEquals("001", tableInstant); + calls.incrementAndGet(); + Map result = new HashMap<>(); + for (String key : keys) { + result.put(key, new HoodieRecordGlobalLocation("p", "001", "f")); + } + return result; + }; + + List result = new RecordIndexVectorCandidateArbiter(lookup, 2) + .arbitrate(HoodieListData.eager(input), request(VectorStalePolicy.FALLBACK), snapshot(), null) + .collectAsList(); + + assertEquals(5, result.size()); + assertEquals(3, calls.get(), "five keys with batch size two require three RLI calls"); + } + + @Test + void allDeletedProducesEmptyOutput() { + List input = new ArrayList<>(); + input.add(candidate("k1", "p", "f", "001")); + input.add(candidate("k2", "p", "f", "001")); + RecordIndexLookup emptyLookup = (keys, tableInstant) -> new HashMap<>(); + HoodieData data = HoodieListData.eager(input); + List result = + new RecordIndexVectorCandidateArbiter(emptyLookup) + .arbitrate(data, request(VectorStalePolicy.FALLBACK), snapshot(), null).collectAsList(); + assertEquals(0, result.size(), "all RLI misses -> all DELETED -> empty"); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestThresholdVectorExecutionModeSelector.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestThresholdVectorExecutionModeSelector.java new file mode 100644 index 0000000000000..2c155a11fd4ed --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestThresholdVectorExecutionModeSelector.java @@ -0,0 +1,79 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.apache.hudi.common.index.vector.VectorDistanceMetric; + +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * Verifies the normative execution-locality rule (RFC-109 §11A) for selector version + * {@code candidate-threshold-v1}: explicit modes are honored, and AUTO selects LOCAL iff + * {@code maxRerankCandidates <= localExecutionThreshold} (boundary inclusive). + */ +public class TestThresholdVectorExecutionModeSelector { + + private final VectorExecutionModeSelector selector = new ThresholdVectorExecutionModeSelector(); + + private VectorSearchRequest requestWith(VectorExecutionMode mode, int maxRerank, int threshold) { + VectorSearchBudget budget = new VectorSearchBudget( + 5000L, 256, 128, maxRerank, Integer.MAX_VALUE, 4, mode, threshold, DeadlinePolicy.FAIL); + return new VectorSearchRequest( + "embedding", new float[] {0.1f, 0.2f}, VectorDistanceMetric.L2, + 10, 32, 50, true, null, budget); + } + + @Test + void autoSelectsLocalAtOrBelowThreshold() { + VectorExecutionDecision below = selector.select(requestWith(VectorExecutionMode.AUTO, 4096, 8192)); + assertEquals(VectorExecutionMode.LOCAL, below.getSelectedMode()); + + // Boundary is inclusive: maxRerank == threshold selects LOCAL. + VectorExecutionDecision boundary = selector.select(requestWith(VectorExecutionMode.AUTO, 8192, 8192)); + assertEquals(VectorExecutionMode.LOCAL, boundary.getSelectedMode()); + } + + @Test + void autoSelectsDistributedAboveThreshold() { + VectorExecutionDecision above = selector.select(requestWith(VectorExecutionMode.AUTO, 8193, 8192)); + assertEquals(VectorExecutionMode.DISTRIBUTED, above.getSelectedMode()); + } + + @Test + void explicitModesAreHonoredRegardlessOfThreshold() { + // Explicit LOCAL even when far above threshold. + VectorExecutionDecision forcedLocal = selector.select(requestWith(VectorExecutionMode.LOCAL, 1_000_000, 8192)); + assertEquals(VectorExecutionMode.LOCAL, forcedLocal.getSelectedMode()); + + // Explicit DISTRIBUTED even when well below threshold. + VectorExecutionDecision forcedDistributed = selector.select(requestWith(VectorExecutionMode.DISTRIBUTED, 8, 8192)); + assertEquals(VectorExecutionMode.DISTRIBUTED, forcedDistributed.getSelectedMode()); + } + + @Test + void decisionRecordsInputsAndSelectorVersion() { + VectorExecutionDecision d = selector.select(requestWith(VectorExecutionMode.AUTO, 4096, 8192)); + assertEquals(VectorExecutionMode.AUTO, d.getRequestedMode()); + assertEquals(4096, d.getMaxRerankCandidates()); + assertEquals(8192, d.getLocalExecutionThreshold()); + assertEquals(VectorExecutionModeSelector.SELECTOR_VERSION, d.getSelectorVersion()); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestVectorCandidateOverlay.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestVectorCandidateOverlay.java new file mode 100644 index 0000000000000..22e883e1cab4b --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestVectorCandidateOverlay.java @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file for details. + */ + +package org.apache.hudi.common.index.vector.search; + +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class TestVectorCandidateOverlay { + + private static VectorCandidate candidate(String key, double distance, String fileId) { + return new VectorCandidate(key, 1, 0, distance, + new VectorPostingLocator(1, 1, 0, 0, 0, "p", fileId, "001", 1)); + } + + @Test + void deltaReplacesBaseEvenWhenDeltaDistanceIsWorse() { + List result = VectorCandidateOverlay.resolve( + Arrays.asList(candidate("updated", 1, "old"), candidate("other", 3, "f")), + Collections.singletonList(candidate("updated", 9, "new")), + Collections.emptySet(), 2, 1); + + assertEquals(Arrays.asList("other", "updated"), keys(result)); + assertEquals("new", result.get(1).getPostingLocator().getFileId()); + } + + @Test + void tombstoneSuppressesBothBaseAndDeltaCopies() { + List result = VectorCandidateOverlay.resolve( + Arrays.asList(candidate("deleted", 1, "old"), candidate("live", 2, "f")), + Collections.singletonList(candidate("deleted", 0.5, "new")), + Collections.singleton("deleted"), 2, 1); + + assertEquals(Collections.singletonList("live"), keys(result)); + } + + @Test + void overlaySlackBackfillsSuppressedBaseFinalist() { + List base = Arrays.asList( + candidate("deleted", 1, "f"), candidate("second", 2, "f"), candidate("backfill", 3, "f")); + + List withoutSlack = VectorCandidateOverlay.resolve( + base, Collections.emptyList(), Collections.singleton("deleted"), 2, 0); + List withSlack = VectorCandidateOverlay.resolve( + base, Collections.emptyList(), Collections.singleton("deleted"), 2, 1); + + assertEquals(Collections.singletonList("second"), keys(withoutSlack)); + assertEquals(Arrays.asList("second", "backfill"), keys(withSlack)); + } + + private static List keys(List candidates) { + return candidates.stream().map(VectorCandidate::getRecordKey).collect(Collectors.toList()); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestVectorContinuation.java b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestVectorContinuation.java new file mode 100644 index 0000000000000..d91b6fcdb3e43 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/common/index/vector/search/TestVectorContinuation.java @@ -0,0 +1,130 @@ +/* + * 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.hudi.common.index.vector.search; + +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Verifies the continuation core (RFC-109 §10): windowing the retained ordered pool into batches + * without rescanning, and accumulating K live exact results with dedup. Simulates stale/deleted + * crowding that forces the reranker to draw additional batches. + */ +public class TestVectorContinuation { + + private static VectorRowRequest row(String key, double dist, VectorCandidateState state) { + return new VectorRowRequest(key, 0L, state, dist); + } + + /** Drives the continuation loop exactly as the reranker will: draw while needsMore && hasMore. */ + private static VectorTopKAccumulator drive(List pool, int topK, + int initial, int batch, int max, int[] outConsumedBatches) { + VectorTopKAccumulator acc = new VectorTopKAccumulator(topK); + VectorContinuationController ctl = + new VectorContinuationController<>(pool, initial, batch, max); + int batches = 0; + while (acc.needsMore() && ctl.hasMore()) { + List b = ctl.nextBatch(); + batches++; + for (VectorRowRequest r : b) { + if (r.getState() != VectorCandidateState.DELETED) { // live (SERVE or STALE-fetched-live) + acc.offer(r.getRecordKey(), r.getApproximateDistance(), null); + } + } + } + if (outConsumedBatches != null) { + outConsumedBatches[0] = batches; + outConsumedBatches[1] = ctl.consumed(); + } + return acc; + } + + @Test + void staleDeletedCrowdingTriggersContinuationUntilKLive() { + List pool = new ArrayList<>(); + // rows 0-2 DELETED, row 3 STALE(live); rows 4-19 SERVE(live). Distances = index. + pool.add(row("r0", 0, VectorCandidateState.DELETED)); + pool.add(row("r1", 1, VectorCandidateState.DELETED)); + pool.add(row("r2", 2, VectorCandidateState.DELETED)); + pool.add(row("r3", 3, VectorCandidateState.STALE)); + for (int i = 4; i < 20; i++) { + pool.add(row("r" + i, i, VectorCandidateState.SERVE)); + } + + int[] out = new int[2]; + VectorTopKAccumulator acc = drive(pool, 5, 4, 4, 4096, out); + + assertEquals(5, acc.liveCount(), "must accumulate exactly K live despite crowding"); + assertFalse(acc.needsMore()); + // batch1 (rows 0-3) -> 1 live; batch2 (rows 4-7) -> 4 live => 5. Two batches, consumed 8. + assertEquals(2, out[0], "should have drawn 2 batches"); + assertEquals(8, out[1], "consumed 8 retained candidates (no rescan)"); + + List top = acc.topK(); + assertEquals(5, top.size()); + assertEquals("r3", top.get(0).getRecordKey(), "nearest live is the STALE-but-live r3"); + assertEquals("r7", top.get(4).getRecordKey()); + } + + @Test + void exhaustsPoolAndReturnsPartialWhenTooFewLive() { + List pool = new ArrayList<>(); + pool.add(row("a", 0, VectorCandidateState.SERVE)); + pool.add(row("b", 1, VectorCandidateState.DELETED)); + pool.add(row("c", 2, VectorCandidateState.SERVE)); + pool.add(row("d", 3, VectorCandidateState.DELETED)); + pool.add(row("e", 4, VectorCandidateState.SERVE)); + + int[] out = new int[2]; + VectorTopKAccumulator acc = drive(pool, 5, 4, 4, 4096, out); + + assertEquals(3, acc.liveCount(), "only 3 live exist -> partial"); + assertTrue(acc.needsMore(), "still needs more but pool is exhausted"); + assertEquals(5, out[1], "consumed the whole retained pool, no more"); + } + + @Test + void accumulatorDedupsByRecordKeyKeepingSmallerDistance() { + VectorTopKAccumulator acc = new VectorTopKAccumulator(3); + acc.offer("x", 5.0, null); + acc.offer("x", 2.0, null); // same key, smaller distance wins + acc.offer("y", 3.0, null); + assertEquals(2, acc.liveCount(), "duplicate key counted once"); + assertEquals("x", acc.topK().get(0).getRecordKey()); + assertEquals(2.0, acc.topK().get(0).getDistance(), 1e-9); + } + + @Test + void maxRerankCandidatesCapsTheRetainedWindow() { + List pool = new ArrayList<>(); + for (int i = 0; i < 100; i++) { + pool.add(row("r" + i, i, VectorCandidateState.DELETED)); // all deleted -> never satisfies K + } + int[] out = new int[2]; + VectorTopKAccumulator acc = drive(pool, 5, 4, 4, 12, out); // cap at 12 + assertEquals(0, acc.liveCount()); + assertEquals(12, out[1], "must not draw beyond maxRerankCandidates"); + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/common/serialization/TestBufferedRecordSerializer.java b/hudi-common/src/test/java/org/apache/hudi/common/serialization/TestBufferedRecordSerializer.java index 0f189be7b6d86..b8892f23902ab 100644 --- a/hudi-common/src/test/java/org/apache/hudi/common/serialization/TestBufferedRecordSerializer.java +++ b/hudi-common/src/test/java/org/apache/hudi/common/serialization/TestBufferedRecordSerializer.java @@ -54,7 +54,7 @@ void testAvroRecordSerAndDe() throws IOException { Assertions.assertEquals(record, result); avroRecordSerializer = new AvroRecordSerializer(integer -> HoodieMetadataRecord.SCHEMA$); - HoodieMetadataRecord metadataRecord = new HoodieMetadataRecord("__all_partitions__", 1, new HashMap<>(), null, null, null, null); + HoodieMetadataRecord metadataRecord = new HoodieMetadataRecord("__all_partitions__", 1, new HashMap<>(), null, null, null, null, null); avroBytes = avroRecordSerializer.serialize(metadataRecord); result = avroRecordSerializer.deserialize(avroBytes, 1); for (int i = 0; i < metadataRecord.getSchema().getFields().size(); i++) { @@ -80,7 +80,7 @@ void testBufferedRecordSerAndDe() throws IOException { avroRecordSerializer = new AvroRecordSerializer(integer -> HoodieMetadataRecord.SCHEMA$); bufferedRecordSerializer = new BufferedRecordSerializer<>(avroRecordSerializer); - HoodieMetadataRecord metadataRecord = new HoodieMetadataRecord("__all_partitions__", 1, new HashMap<>(), null, null, null, null); + HoodieMetadataRecord metadataRecord = new HoodieMetadataRecord("__all_partitions__", 1, new HashMap<>(), null, null, null, null, null); bufferedRecord = new BufferedRecord<>("__all_partitions__", 0, metadataRecord, 1, null); bytes = bufferedRecordSerializer.serialize(bufferedRecord); result = bufferedRecordSerializer.deserialize(bytes); diff --git a/hudi-common/src/test/java/org/apache/hudi/metadata/TestMetadataPartitionType.java b/hudi-common/src/test/java/org/apache/hudi/metadata/TestMetadataPartitionType.java index 8aa5879794c84..7f2a47f291801 100644 --- a/hudi-common/src/test/java/org/apache/hudi/metadata/TestMetadataPartitionType.java +++ b/hudi-common/src/test/java/org/apache/hudi/metadata/TestMetadataPartitionType.java @@ -116,10 +116,14 @@ public void testPartitionEnabledByConfigOnly(MetadataPartitionType partitionType } List enabledPartitions = MetadataPartitionType.getEnabledPartitions(metadataConfigBuilder.build(), metaClient); - // Verify partition type is enabled due to config assertEquals(expectedEnabledPartitions, enabledPartitions.size()); Set validPartitions = Arrays.stream(MetadataPartitionType.getValidValues(tableVersion)).collect(Collectors.toSet()); - assertTrue(!validPartitions.contains(partitionType) || enabledPartitions.contains(partitionType) || MetadataPartitionType.ALL_PARTITIONS.equals(partitionType)); + if (partitionType == MetadataPartitionType.VECTOR_INDEX) { + // Vector indexes are enabled by a persisted CREATE INDEX definition, never by metadata config alone. + assertFalse(enabledPartitions.contains(partitionType)); + } else { + assertTrue(!validPartitions.contains(partitionType) || enabledPartitions.contains(partitionType) || MetadataPartitionType.ALL_PARTITIONS.equals(partitionType)); + } } @Test @@ -205,6 +209,7 @@ public void testFromPartitionPath() { assertEquals(MetadataPartitionType.BLOOM_FILTERS, MetadataPartitionType.fromPartitionPath("bloom_filters")); assertEquals(MetadataPartitionType.RECORD_INDEX, MetadataPartitionType.fromPartitionPath("record_index")); assertEquals(MetadataPartitionType.PARTITION_STATS, MetadataPartitionType.fromPartitionPath("partition_stats")); + assertEquals(MetadataPartitionType.VECTOR_INDEX, MetadataPartitionType.fromPartitionPath("vector_index_embedding")); assertThrows(IllegalArgumentException.class, () -> MetadataPartitionType.fromPartitionPath("unknown")); } @@ -217,6 +222,7 @@ public void testGetMetadataPartitionRecordType() { assertEquals(5, MetadataPartitionType.RECORD_INDEX.getRecordType()); assertEquals(6, MetadataPartitionType.PARTITION_STATS.getRecordType()); assertEquals(7, MetadataPartitionType.SECONDARY_INDEX.getRecordType()); + assertEquals(8, MetadataPartitionType.VECTOR_INDEX.getRecordType()); } @ParameterizedTest @@ -225,7 +231,8 @@ public void testGetNonExpressionIndexPath(MetadataPartitionType partitionType) { HoodieTableMetaClient metaClient = mock(HoodieTableMetaClient.class); String expressionIndexName = "expr_index_dummyExpressionIndex"; String secondaryIndexName = "secondary_index_dummySecondaryIndex"; - HoodieIndexMetadata indexMetadata = getIndexMetadata(expressionIndexName, secondaryIndexName); + String vectorIndexName = "vector_index_dummyVectorIndex"; + HoodieIndexMetadata indexMetadata = getIndexMetadata(expressionIndexName, secondaryIndexName, vectorIndexName); when(metaClient.getIndexMetadata()).thenReturn(Option.of(indexMetadata)); // Mock getIndexForMetadataPartition for both index names @@ -233,17 +240,21 @@ public void testGetNonExpressionIndexPath(MetadataPartitionType partitionType) { .thenReturn(Option.of(indexMetadata.getIndexDefinitions().get(expressionIndexName))); when(metaClient.getIndexForMetadataPartition(secondaryIndexName)) .thenReturn(Option.of(indexMetadata.getIndexDefinitions().get(secondaryIndexName))); - + when(metaClient.getIndexForMetadataPartition(vectorIndexName)) + .thenReturn(Option.of(indexMetadata.getIndexDefinitions().get(vectorIndexName))); + if (partitionType == MetadataPartitionType.EXPRESSION_INDEX) { assertEquals(expressionIndexName, partitionType.getPartitionPath(metaClient, expressionIndexName)); } else if (partitionType == MetadataPartitionType.SECONDARY_INDEX) { assertEquals(secondaryIndexName, partitionType.getPartitionPath(metaClient, secondaryIndexName)); + } else if (partitionType == MetadataPartitionType.VECTOR_INDEX) { + assertEquals(vectorIndexName, partitionType.getPartitionPath(metaClient, vectorIndexName)); } else { assertEquals(partitionType.getPartitionPath(), partitionType.getPartitionPath(metaClient, null)); } } - private static HoodieIndexMetadata getIndexMetadata(String expressionIndexName, String secondaryIndexName) { + private static HoodieIndexMetadata getIndexMetadata(String expressionIndexName, String secondaryIndexName, String vectorIndexName) { Map indexDefinitions = new HashMap<>(); HoodieIndexDefinition expressionIndexDefinition = HoodieIndexDefinition.newBuilder() .withIndexName(expressionIndexName) @@ -261,6 +272,13 @@ private static HoodieIndexMetadata getIndexMetadata(String expressionIndexName, .withSourceFields(Collections.singletonList("name")) .build(); indexDefinitions.put(secondaryIndexName, secondaryIndexDefinition); + HoodieIndexDefinition vectorIndexDefinition = HoodieIndexDefinition.newBuilder() + .withIndexName(vectorIndexName) + .withIndexType("VECTOR") + .withVersion(HoodieIndexVersion.getCurrentVersion(HoodieTableVersion.current(), vectorIndexName)) + .withSourceFields(Collections.singletonList("embedding")) + .build(); + indexDefinitions.put(vectorIndexName, vectorIndexDefinition); return new HoodieIndexMetadata(indexDefinitions); } @@ -278,7 +296,9 @@ public void testExceptionForMissingExpressionIndexMetadata() { @Test public void testIndexNameWithoutPrefix() { for (MetadataPartitionType partitionType : MetadataPartitionType.getValidValues(HoodieTableVersion.current())) { - String userIndexName = MetadataPartitionType.isExpressionOrSecondaryIndex(partitionType.getPartitionPath()) ? "idx" : ""; + boolean hasDynamicPartitionPath = MetadataPartitionType.isExpressionOrSecondaryIndex(partitionType.getPartitionPath()) + || partitionType == MetadataPartitionType.VECTOR_INDEX; + String userIndexName = hasDynamicPartitionPath ? "idx" : ""; HoodieIndexDefinition indexDefinition = createIndexDefinition(partitionType, userIndexName, partitionType.name(), null, null, null); assertEquals(partitionType.getIndexNameWithoutPrefix(indexDefinition), userIndexName); } diff --git a/hudi-common/src/test/java/org/apache/hudi/metadata/TestVectorIndexMetadataKey.java b/hudi-common/src/test/java/org/apache/hudi/metadata/TestVectorIndexMetadataKey.java new file mode 100644 index 0000000000000..007125cd3779b --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/metadata/TestVectorIndexMetadataKey.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.hudi.metadata; + +import org.junit.jupiter.api.Test; + +import java.nio.charset.StandardCharsets; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class TestVectorIndexMetadataKey { + + @Test + void testBlockReservedAndDeltaRangesSortInSinglePrefixScan() { + assertOrdered( + VectorIndexMetadataKey.postingBlock(1, 7, 0, 0x00000000L), + VectorIndexMetadataKey.postingBlock(1, 7, 0, 0x00000001L), + VectorIndexMetadataKey.postingBlock(1, 7, 0, VectorIndexMetadataKey.MAX_PACKED_BLOCK_ID), + VectorIndexMetadataKey.postingBlock(1, 7, 0, VectorIndexMetadataKey.FIRST_RESERVED_BLOCK_ID), + VectorIndexMetadataKey.postingBlock(1, 7, 0, VectorIndexMetadataKey.LAST_RESERVED_BLOCK_ID), + VectorIndexMetadataKey.postingDelta(1, 7, 0, "any")); + } + + @Test + void testUnsignedComponentBoundaries() { + assertOrdered( + VectorIndexMetadataKey.postingBlock(1, 0x7FFFFFFF, 0, 0), + VectorIndexMetadataKey.postingBlock(1, 0x80000000, 0, 0), + VectorIndexMetadataKey.postingBlock(1, 0xFFFFFFFF, 0, 0)); + assertOrdered( + VectorIndexMetadataKey.postingBlock(1, 7, 0x7FFF, 0), + VectorIndexMetadataKey.postingBlock(1, 7, 0x8000, 0)); + assertOrdered( + VectorIndexMetadataKey.postingBlock(1, 7, 0, 0x7FFFFFFFL), + VectorIndexMetadataKey.postingBlock(1, 7, 0, 0x80000000L)); + } + + @Test + void testBigEndianCarryBoundaries() { + assertOrdered( + VectorIndexMetadataKey.postingBlock(1, 0x000000FF, 0xFFFF, VectorIndexMetadataKey.DELTA_BLOCK_ID), + VectorIndexMetadataKey.postingBlock(1, 0x00000100, 0x0000, 0)); + assertOrdered( + VectorIndexMetadataKey.postingBlock(1, 7, 0x00FF, VectorIndexMetadataKey.DELTA_BLOCK_ID), + VectorIndexMetadataKey.postingBlock(1, 7, 0x0100, 0)); + } + + @Test + void testDeltaKeyTailUtf8OrderingAndDecode() { + String a = VectorIndexMetadataKey.postingDelta(1, 7, 0, "a"); + String aa = VectorIndexMetadataKey.postingDelta(1, 7, 0, "aa"); + String b = VectorIndexMetadataKey.postingDelta(1, 7, 0, "b"); + String eAcute = VectorIndexMetadataKey.postingDelta(1, 7, 0, "é"); + + assertOrdered(a, aa, b, eAcute); + assertEquals("é", VectorIndexMetadataKey.postingRecordKey(eAcute)); + assertEquals(7, VectorIndexMetadataKey.postingClusterId(eAcute)); + assertEquals(0, VectorIndexMetadataKey.postingShard(eAcute)); + } + + @Test + void testFamilyMajorOrdering() { + assertOrdered( + VectorIndexMetadataKey.activeManifest(), + VectorIndexMetadataKey.manifest(1), + VectorIndexMetadataKey.manifest(2), + VectorIndexMetadataKey.quantizer(1, 0), + VectorIndexMetadataKey.centroids(1, 0), + VectorIndexMetadataKey.clusterStats(1, 7), + VectorIndexMetadataKey.postingBlock(1, 7, 0, 0), + VectorIndexMetadataKey.sourceInstantMarker(1, "20260724010101")); + } + + @Test + void testSourceInstantMarkersSortWithinGenerationAndDecode() { + String prefix = VectorIndexMetadataKey.sourceInstantMarkerPrefix(3); + String first = VectorIndexMetadataKey.sourceInstantMarker(3, "20260724010101"); + String second = VectorIndexMetadataKey.sourceInstantMarker(3, "20260724010202"); + String nextGeneration = VectorIndexMetadataKey.sourceInstantMarker(4, "20260724000000"); + + assertOrdered(prefix, first, second, nextGeneration); + assertEquals("20260724010101", VectorIndexMetadataKey.sourceInstant(first)); + assertNull(VectorIndexMetadataKey.sourceInstant(VectorIndexMetadataKey.manifest(3))); + assertNull(VectorIndexMetadataKey.sourceInstant(prefix)); + assertThrows(IllegalArgumentException.class, + () -> VectorIndexMetadataKey.sourceInstantMarker(3, "")); + assertThrows(IllegalArgumentException.class, + () -> VectorIndexMetadataKey.sourceInstantMarker(3, null)); + assertTrue(VectorIndexMetadataKey.compareUnsigned( + second, VectorIndexMetadataKey.exclusiveEnd(prefix)) < 0); + } + + @Test + void testPrefixScanExclusiveEnd() { + String prefix = VectorIndexMetadataKey.postingPrefix(1, 7, 0); + String end = VectorIndexMetadataKey.exclusiveEnd(prefix); + + assertTrue(VectorIndexMetadataKey.compareUnsigned(prefix, end) < 0); + assertTrue(VectorIndexMetadataKey.compareUnsigned( + VectorIndexMetadataKey.postingDelta(1, 7, 0, "z"), end) < 0); + assertTrue(VectorIndexMetadataKey.compareUnsigned( + VectorIndexMetadataKey.postingBlock(1, 7, 1, 0), end) >= 0); + assertNull(VectorIndexMetadataKey.exclusiveEnd(new String(new byte[] {(byte) 0xFF}, StandardCharsets.ISO_8859_1))); + } + + private static void assertOrdered(String... keys) { + for (int i = 1; i < keys.length; i++) { + assertTrue(VectorIndexMetadataKey.compareUnsigned(keys[i - 1], keys[i]) < 0, + "key " + (i - 1) + " should sort before key " + i); + } + } +} diff --git a/hudi-common/src/test/java/org/apache/hudi/metadata/TestVectorIndexMetadataPayload.java b/hudi-common/src/test/java/org/apache/hudi/metadata/TestVectorIndexMetadataPayload.java new file mode 100644 index 0000000000000..84e32a8936b56 --- /dev/null +++ b/hudi-common/src/test/java/org/apache/hudi/metadata/TestVectorIndexMetadataPayload.java @@ -0,0 +1,162 @@ +/* + * 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.hudi.metadata; + +import org.apache.hudi.avro.model.HoodieVectorIndexActiveManifest; +import org.apache.hudi.avro.model.HoodieVectorIndexCentroids; +import org.apache.hudi.avro.model.HoodieVectorIndexClusterStats; +import org.apache.hudi.avro.model.HoodieVectorIndexManifest; +import org.apache.hudi.avro.model.HoodieVectorIndexPostingDelta; +import org.apache.hudi.avro.model.HoodieVectorIndexQuantizer; +import org.apache.hudi.avro.model.HoodieVectorIndexSourceInstantMarker; +import org.apache.hudi.common.model.HoodieRecord; + +import org.junit.jupiter.api.Test; + +import java.nio.ByteBuffer; +import java.util.Arrays; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class TestVectorIndexMetadataPayload { + + @Test + void testPostingRecordCarriesCanonicalLookupMetadata() { + HoodieRecord record = HoodieMetadataPayload.createVectorIndexPostingRecord( + 7, + "rk-1", + 3, + 1, + "file-group-1", + "dt=2026-04-01", + "20260603120000", + new byte[] {0x01, 0x02}, + 1.5f, + 123456789L, + "vector_index_demo"); + + assertTrue(record.getData().getVectorIndexMetadata().isPresent()); + HoodieVectorIndexPostingDelta delta = + (HoodieVectorIndexPostingDelta) record.getData().getVectorIndexMetadata().get(); + assertEquals("rk-1", delta.getRecordKey()); + assertEquals(3, VectorIndexMetadataKey.postingClusterId(record.getRecordKey())); + assertEquals(1, VectorIndexMetadataKey.postingShard(record.getRecordKey())); + assertEquals("file-group-1", delta.getFileGroupId()); + assertEquals("dt=2026-04-01", delta.getPartitionPath()); + assertEquals("20260603120000", delta.getBaseInstantTime()); + } + + @Test + void testActiveManifestCarriesReaderVisibleGeneration() { + HoodieRecord record = HoodieMetadataPayload.createVectorIndexActiveManifestRecord( + 2, "vector_index_demo"); + + HoodieVectorIndexActiveManifest manifest = + (HoodieVectorIndexActiveManifest) record.getData().getVectorIndexMetadata().get(); + assertEquals(VectorIndexMetadataKey.activeManifest(), record.getRecordKey()); + assertEquals(1, manifest.getIndexVersion()); + assertEquals(2, manifest.getActiveGeneration()); + } + + @Test + void testEpochFreeCentroidRecordUsesGenerationAndChunkKey() { + HoodieRecord record = HoodieMetadataPayload.createVectorIndexCentroidsRecord( + 2, + 7, + ByteBuffer.wrap(new byte[] {1, 0, 0, 0}), + ByteBuffer.wrap(new byte[] {2, 3}), + ByteBuffer.wrap(new byte[] {4, 5}), + "vector_index_demo"); + + HoodieVectorIndexCentroids centroids = + (HoodieVectorIndexCentroids) record.getData().getVectorIndexMetadata().get(); + assertEquals(VectorIndexMetadataKey.centroids(2, 7), record.getRecordKey()); + assertEquals(2, centroids.getCentroidBytes().remaining()); + assertNull(HoodieVectorIndexCentroids.getClassSchema().getField("centroidEpoch")); + } + + @Test + void testManifestCarriesVerifiedContiguousFrontierWithoutEpoch() { + HoodieRecord record = HoodieMetadataPayload.createVectorIndexManifestRecord( + 2, "build-2", "BUILDING", 128, 128, 16, 2, 1, 64, 1, 8, + "COSINE", true, true, "embedding", 524288, 2048, + 1, 1, 1.9, 1.0e-3, 1.0, 1.0e-3, 4, "sha256:centroids", + 4096, 1024, "20260724000000", "20260724010101", 123L, "vector_index_demo"); + + HoodieVectorIndexManifest manifest = + (HoodieVectorIndexManifest) record.getData().getVectorIndexMetadata().get(); + assertEquals("20260724000000", manifest.getBootstrapInstant()); + assertEquals("20260724010101", manifest.getVerifiedFrontier()); + assertEquals(8, manifest.getFileGroupCount()); + assertEquals(1, manifest.getBlockFormatVersion()); + assertEquals(1, manifest.getFactorVersion()); + assertEquals(1.9, manifest.getKappa()); + assertEquals(1.0e-3, manifest.getGMin()); + assertEquals(1.0, manifest.getEps1Max()); + assertEquals(1.0e-3, manifest.getEpsNRel()); + assertEquals(4, manifest.getCentroidChunkCount()); + assertEquals("sha256:centroids", manifest.getCentroidChecksum()); + assertEquals("BUILDING", manifest.getState().toString()); + assertNull(HoodieVectorIndexManifest.getClassSchema().getField("centroidEpoch")); + } + + @Test + void testClusterManifestPersistsRoutingFields() { + HoodieRecord record = HoodieMetadataPayload.createVectorIndexClusterManifestRecord( + 2, 9, 2, Arrays.asList("fg-a", "fg-b"), 17L, 123L, "vector_index_demo"); + + HoodieVectorIndexClusterStats stats = + (HoodieVectorIndexClusterStats) record.getData().getVectorIndexMetadata().get(); + assertEquals(0, stats.getRoutingVersion()); + assertEquals(2, stats.getShardCount()); + assertEquals(Arrays.asList("fg-a", "fg-b"), stats.getFileGroupIds()); + assertEquals(17L, stats.getLiveCount()); + assertNull(HoodieVectorIndexClusterStats.getClassSchema().getField("centroidEpoch")); + } + + @Test + void testSourceInstantMarkerCarriesSourceIdentity() { + HoodieRecord record = HoodieMetadataPayload.createVectorIndexSourceInstantMarkerRecord( + 2, "20260724010101", "vector_index_demo"); + + HoodieVectorIndexSourceInstantMarker marker = + (HoodieVectorIndexSourceInstantMarker) record.getData().getVectorIndexMetadata().get(); + assertEquals(VectorIndexMetadataKey.sourceInstantMarker(2, "20260724010101"), record.getRecordKey()); + assertEquals("20260724010101", marker.getDataInstant().toString()); + } + + @Test + void testQuantizerMetadataRecordCarriesRaBitQConfig() { + HoodieRecord record = HoodieMetadataPayload.createVectorIndexQuantizerMetadataRecord( + "IVF_RABITQ", + 96, + 42L, + true, + "vector_index_demo"); + + assertTrue(record.getData().getVectorIndexMetadata().isPresent()); + HoodieVectorIndexQuantizer quantizer = + (HoodieVectorIndexQuantizer) record.getData().getVectorIndexMetadata().get(); + assertEquals("IVF_RABITQ", quantizer.getQuantizerType()); + assertEquals(42L, quantizer.getRandomSeed()); + } +}