chrevanthreddy commented on code in PR #20107: URL: https://github.com/apache/hudi/pull/20107#discussion_r4384297540
########## hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkVectorIndexBootstrap.java: ########## @@ -0,0 +1,531 @@ +/* + * 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.common.data.HoodieData; +import org.apache.hudi.common.index.vector.PostingBlockBuilder; +import org.apache.hudi.common.index.vector.RaBitQEncoder; +import org.apache.hudi.common.index.vector.VectorDistanceMetric; +import org.apache.hudi.common.index.vector.VectorIndexBootstrapUtils; +import org.apache.hudi.common.index.vector.VectorIndexOptions; +import org.apache.hudi.common.model.HoodieIndexDefinition; +import org.apache.hudi.common.model.HoodieRecord; +import org.apache.hudi.common.schema.HoodieSchema; +import org.apache.hudi.data.HoodieJavaRDD; +import org.apache.hudi.spark.index.vector.TwoLevelKMeansBootstrap$; + +import lombok.extern.slf4j.Slf4j; +import org.apache.spark.api.java.JavaPairRDD; +import org.apache.spark.api.java.JavaRDD; +import org.apache.spark.api.java.JavaSparkContext; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.sql.SparkSession; + +import java.io.Serializable; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import scala.Tuple2; + +/** + * Two-phase vector index bootstrap for Spark. + * + * <p>Phase 1: Sample vectors → train centroids with Spark ML KMeans. Training + * still uses a bounded sample (FAISS, ScaNN, and Milvus all use sample-based + * centroid training for IVF), but the iterative KMeans work stays distributed + * instead of collecting the sample to the driver. + * + * <p>Phase 2: Broadcast centroids → single-pass {@code mapPartitions} that + * assigns each vector to its nearest centroid, computes RaBitQ binary codes, + * and emits posting metadata records with per-candidate locators. Cluster + * statistics (populations, shard counts) are gathered via a lightweight + * {@code countByKey} before the main encoding pass. + * + * <p>Compared to the previous driver-local KMeans approach: + * <ul> + * <li>No reflection — direct {@link HoodieMetadataPayload} method calls</li> + * <li>Proper {@code unpersist()} in try/finally</li> + * <li>generation ordinal allocated from existing manifest keys</li> + * </ul> + */ +@Slf4j +public final class SparkVectorIndexBootstrap { + + /** Hard cap for percentile-driven training samples. */ + private static final int MAX_PERCENTILE_TRAINING_SAMPLE = 10_000_000; + + /** Absolute minimum sample floor for production index quality. */ + private static final int MIN_TRAINING_SAMPLE = 1_000_000; + + /** Minimum sample multiplier per cluster. FAISS recommends 256×K minimum. */ + private static final int MIN_SAMPLE_PER_CLUSTER = 256; + + /** Lower bound for percentile-driven sampling: 0.5% of total vectors. */ + private static final double MIN_SAMPLE_PERCENT = 0.005d; + + /** Upper bound for percentile-driven sampling: 1% of total vectors. */ + private static final double MAX_SAMPLE_PERCENT = 0.01d; + + private SparkVectorIndexBootstrap() { + } + + /** + * Run the two-phase vector index bootstrap and return all MDT records. + * + * @param jsc Spark context + * @param vectorRows RDD of (recordKey, partitionPath, fileId, vectorBytes) — already + * read from base-table file slices + * @param indexDef the index definition + * @param vectorType resolved vector element type from the table schema + * @param generation vector generation ordinal + * @param lastUpdatedTs timestamp for MDT bookkeeping + * @return all metadata records to write into the vector index MDT partition + */ + public static HoodieData<HoodieRecord> bootstrap(JavaSparkContext jsc, + JavaRDD<VectorRow> vectorRows, + HoodieIndexDefinition indexDef, + HoodieSchema.Vector.VectorElementType vectorType, + int dimension, + int generation, + long lastUpdatedTs) { + Map<String, String> options = indexDef.getIndexOptions(); + String indexName = indexDef.getIndexName(); + String vectorColumn = indexDef.getSourceFields().isEmpty() ? "" : indexDef.getSourceFields().get(0); + VectorIndexOptions.ResolvedOptions resolvedOptions = VectorIndexOptions.resolve(options); + int numClusters = resolvedOptions.numClusters; + VectorDistanceMetric metric = resolvedOptions.metric; + + String quantizerType = resolvedOptions.quantizer.name(); + long quantizerSeed = resolvedOptions.rabitqSeed; + int rabitqBits = resolvedOptions.rabitqBits; + boolean assumeNormalized = resolvedOptions.assumeNormalized; + boolean storeInMdt = true; + int targetRowsPerShard = 100_000; Review Comment: **Fixed** `3cde5283ea`. Flag and dead branches removed. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
