chrevanthreddy commented on code in PR #20107: URL: https://github.com/apache/hudi/pull/20107#discussion_r4384291016
########## hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexRliArbitrator.java: ########## @@ -0,0 +1,230 @@ +/* + * 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.data.HoodieData; +import org.apache.hudi.common.data.HoodieListData; +import org.apache.hudi.common.model.HoodieRecordGlobalLocation; +import org.apache.hudi.common.util.Option; +import org.apache.hudi.metadata.HoodieTableMetadata; + +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.Objects; +import java.util.Set; +import java.util.stream.Collectors; + +/** Applies the RFC-109 record-level-index arbitration contract to vector finalists. */ +final class VectorIndexRliArbitrator { + + private VectorIndexRliArbitrator() { + } + + /** + * The RFC-109 RLI finalist arbiter. Resolves each finalist's current location from the + * record-level index (one batched {@code readRecordIndexLocationsWithKeys} over the distinct + * finalist keys) and tags it with a {@link VectorIndexArbiter.Decision} plus the resolved + * location. + * + * <p>Unlike {@code VectorIndexMdtSearchUtils.attachRecordLocations}, this does <em>not</em> drop candidates: it tags all + * of them so callers can tally {@code arbiterExclusions.stale} / {@code .deleted} and apply the + * mode-specific action (approx: exclude STALE + DELETED; exact: key-fallback STALE, exclude + * DELETED, positional SERVE). Resolved location semantics: + * + * <ul> + * <li>{@code SERVE}: the posting's own location when present (positional trust), else the RLI + * location.</li> + * <li>{@code STALE}: the RLI current location, so exact mode can key-fetch at the live slice.</li> + * <li>{@code DELETED}: {@code null}.</li> + * </ul> + */ + static HoodieData<ScoredVectorPostingMatch> arbitrateFinalists(HoodieTableMetadata metadataTable, + HoodieData<ScoredVectorPostingMatch> finalists) { + return arbitrateFinalists(metadataTable, finalists, false); + } + + static HoodieData<ScoredVectorPostingMatch> arbitrateFinalists( + HoodieTableMetadata metadataTable, + HoodieData<ScoredVectorPostingMatch> finalists, + boolean partitionedRecordIndex) { + if (!partitionedRecordIndex) { + return arbitrateFinalistsForPartition(metadataTable, finalists, Option.empty()); + } + List<String> partitions = finalists.map(ScoredVectorPostingMatch::getPartitionPath) + .distinct() + .collectAsList(); + HoodieData<ScoredVectorPostingMatch> arbitrated = null; + for (String partition : partitions) { + HoodieData<ScoredVectorPostingMatch> partitionFinalists = finalists + .filter(candidate -> Objects.equals(partition, candidate.getPartitionPath())); + HoodieData<ScoredVectorPostingMatch> partitionResult = arbitrateFinalistsForPartition( + metadataTable, partitionFinalists, Option.ofNullable(partition)); + arbitrated = arbitrated == null ? partitionResult : arbitrated.union(partitionResult); + } + return arbitrated == null ? HoodieListData.eager(Collections.emptyList()) : arbitrated; + } + + private static HoodieData<ScoredVectorPostingMatch> arbitrateFinalistsForPartition( + HoodieTableMetadata metadataTable, + HoodieData<ScoredVectorPostingMatch> finalists, + Option<String> dataTablePartition) { + // Resolve current RLI locations for the finalist keys into a bounded driver-side map, then + // attach per candidate via map(...). The finalist set is a bounded candidate pool and + // {@code finalists} is already persisted upstream, so the two passes are cache hits. + // + // This deliberately avoids leftOuterJoin: HoodiePairData.leftOuterJoin requires both operands to + // share the same backing flavor, but readRecordIndexLocationsWithKeys returns list-backed pair + // data for a single-slice RLI (the common 1-file-group case) and RDD-backed for multi-slice. + // Joining an RDD-backed finalist set against list-backed locations throws ClassCastException. + // Attaching via map(...) preserves the finalists' backing (RDD stays RDD, list stays list). + List<String> distinctKeys = finalists.map(ScoredVectorPostingMatch::getRecordKey) + .distinct() + .collectAsList(); + Map<String, HoodieRecordGlobalLocation> currentLocations = new HashMap<>(); + if (!distinctKeys.isEmpty()) { + metadataTable.readRecordIndexLocationsWithKeys( + HoodieListData.eager(distinctKeys), dataTablePartition) + .collectAsList() + .forEach(pair -> currentLocations.put(pair.getKey(), pair.getValue())); Review Comment: Thanks, agreed: the batch precondition was violated. Finalists are now grouped by data partition and then by `HoodieTableMetadataUtil.mapRecordKeyToFileGroupIndex(key, numFileGroups)` before each `readRecordIndexLocationsWithKeys(keys, Option.of(partition))` call, the same grouping `SparkMetadataTableRecordLevelIndex` / `PartitionedRecordLevelIndexSupport` use. File groups come from `getBucketizedFileGroupsForPartitionedRLI`. A partition with no RLI file groups resolves to an RLI miss. Covered by `TestVectorIndexRliArbitratorPartitionedRli` (keys spread across multiple RLI file groups per partition). ########## hudi-client/hudi-client-common/src/main/java/org/apache/hudi/metadata/index/vector/VectorIndexer.java: ########## @@ -0,0 +1,161 @@ +/* + * 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.index.vector; + +import org.apache.hudi.common.data.HoodieData; +import org.apache.hudi.common.engine.HoodieEngineContext; +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.common.table.HoodieTableMetaClient; +import org.apache.hudi.common.util.ValidationUtils; +import org.apache.hudi.config.HoodieWriteConfig; +import org.apache.hudi.exception.HoodieMetadataException; +import org.apache.hudi.metadata.HoodieTableMetadataUtil; +import org.apache.hudi.metadata.MetadataPartitionType; +import org.apache.hudi.metadata.index.BaseIndexer; +import org.apache.hudi.metadata.index.EngineIndexerSupport; +import org.apache.hudi.metadata.index.model.IndexCleanContext; +import org.apache.hudi.metadata.index.model.IndexInitializationContext; +import org.apache.hudi.metadata.index.model.IndexInitializationPlan; +import org.apache.hudi.metadata.index.model.IndexPartitionAndRecords; +import org.apache.hudi.metadata.index.model.IndexUpdateContext; +import org.apache.hudi.metadata.model.FileSliceAndPartition; + +import lombok.extern.slf4j.Slf4j; + +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.apache.hudi.metadata.MetadataPartitionType.VECTOR_INDEX; + +/** + * Implementation of {@link MetadataPartitionType#VECTOR_INDEX} index. + * <p> + * Owns partition sizing (IVF cluster to file-group mapping) and initialization-plan assembly. + * The engine-specific record generation (reading the source VECTOR column, IVF clustering and + * RaBitQ quantization) is dispatched through {@link EngineIndexerSupport}. + * <p> + * Index update, clean and restore lifecycle is handled in a follow-up change; for now those hooks + * are no-ops so the initial bootstrap path is self-contained and reviewable. + */ +@Slf4j +public class VectorIndexer extends BaseIndexer { + + /** + * Target size per metadata file group used by the sizing guard. Vector postings are routed to + * file groups by cluster id, so this is only a sanity bound against gross over-provisioning. + */ + private static final long TARGET_FILE_GROUP_BYTES = 512L * 1024L * 1024L; + + private final EngineIndexerSupport engineIndexerSupport; + + public VectorIndexer( + HoodieEngineContext engineContext, + HoodieWriteConfig dataTableWriteConfig, + HoodieTableMetaClient dataTableMetaClient, + EngineIndexerSupport engineIndexerSupport) { + super(engineContext, dataTableWriteConfig, dataTableMetaClient); + this.engineIndexerSupport = engineIndexerSupport; + } + + @Override + public List<IndexInitializationPlan> buildInitialization(IndexInitializationContext context) throws IOException { + Set<String> completedPartitions = dataTableMetaClient.getTableConfig().getMetadataPartitions(); + Set<String> vectorIndexPartitionsToInit = dataTableMetaClient.getIndexMetadata() + .map(metadata -> metadata.getIndexDefinitions().values().stream() + .map(HoodieIndexDefinition::getIndexName) + .filter(name -> name.startsWith(VECTOR_INDEX.getPartitionPath())) + .filter(name -> !completedPartitions.contains(name)) + .collect(Collectors.toSet())) + .orElse(Collections.emptySet()); + if (vectorIndexPartitionsToInit.size() != 1) { + if (vectorIndexPartitionsToInit.size() > 1) { + log.warn("Skipping vector index initialization as only one vector index bootstrap at a time " + + "is supported for now. Provided: {}", vectorIndexPartitionsToInit); + } + return Collections.emptyList(); + } + + String indexName = vectorIndexPartitionsToInit.iterator().next(); + HoodieIndexDefinition indexDefinition = HoodieTableMetadataUtil.getHoodieIndexDefinition(indexName, dataTableMetaClient); + ValidationUtils.checkState(indexDefinition != null, "Vector index definition is not present for index " + indexName); + + int fileGroupCount = resolveFileGroupCount(indexDefinition); + + HoodieSchema tableSchema = context.tableSchema().get() + .orElseThrow(() -> new HoodieMetadataException("Table schema is not available for vector index initialization")); + List<FileSliceAndPartition> fileSlices = context.latestFileSlices().get(); + + // Initial bootstrap always allocates generation 1 (no pre-existing MDT partition to advance). + HoodieData<HoodieRecord> records = engineIndexerSupport.generateVectorIndexRecords( + indexDefinition, dataTableMetaClient, fileSlices, tableSchema, 1); + + return Collections.singletonList(IndexInitializationPlan.of(fileGroupCount, indexName, records)); + } + + /** + * Sizes the vector metadata partition by IVF clusters. Posting rows are routed to file groups by + * cluster id ({@code mapVectorPostingKeyToFileGroupIndex}), so + * {@code fileGroupCount = ceil(numClusters / clustersPerFileGroup)} keeps one cluster per file group + * by default. An explicit {@code hoodie.metadata.vector.index.file.group.count} overrides the derived + * value. + */ + private int resolveFileGroupCount(HoodieIndexDefinition indexDefinition) { + VectorIndexOptions.ResolvedOptions resolvedOptions = VectorIndexOptions.resolve(indexDefinition.getIndexOptions()); + int numClusters = resolvedOptions.numClusters; + int clustersPerFileGroup = 1; + int configuredFileGroupCount = dataTableWriteConfig.getMetadataConfig().getVectorIndexFileGroupCount(); + int fileGroupCount = configuredFileGroupCount > 0 + ? configuredFileGroupCount + : Math.max(1, (int) Math.ceil((double) numClusters / clustersPerFileGroup)); + + long targetBlockBytes = 64L * 1024L; + long projectedIndexBytes = Math.max(1L, (long) numClusters * targetBlockBytes); + if ((long) fileGroupCount * TARGET_FILE_GROUP_BYTES > projectedIndexBytes * 10L) { + log.warn("Vector index {} is initializing with {} file groups for projected index size {} bytes. " + + "This exceeds the 10x sizing guard against 512 MiB target file groups; consider setting {} " + + "to a smaller value before bootstrap.", + indexDefinition.getIndexName(), fileGroupCount, projectedIndexBytes, + org.apache.hudi.common.config.HoodieMetadataConfig.VECTOR_INDEX_FILE_GROUP_COUNT.key()); + } + log.info("Initializing vector index {} with {} file groups " + + "[numClusters={}, clustersPerFileGroup={}, configuredFileGroupCount={}]", + indexDefinition.getIndexName(), fileGroupCount, numClusters, clustersPerFileGroup, configuredFileGroupCount); + return fileGroupCount; + } + + @Override + public List<IndexPartitionAndRecords> buildUpdate(IndexUpdateContext context) { + // Vector index incremental maintenance (append/delete postings, generation advance) is handled + // in a follow-up change. Initialization is a full re-bootstrap until then. + return Collections.emptyList(); + } + + @Override + public List<IndexPartitionAndRecords> buildClean(IndexCleanContext context) { + // See buildUpdate: lifecycle maintenance is handled in a follow-up change. + return Collections.emptyList(); + } Review Comment: Yes. This PR is intentionally bootstrap + read path only. Synchronous maintenance lives in the stacked PR **https://github.com/chrevanthreddy/hudi/pull/5** (based on this one): `buildUpdate` encodes inserts, vector updates, deletes and cluster moves per source commit, and skips posting work on compaction (as RLI/SI do). It's validated by an insert / upsert / replay / delete / clustering lifecycle that checks exact search == brute force at 10M (COW + MOR) and at 1B (COW). `buildClean` stays a no-op by design: cleaning removes no live rows, so postings don't change. Until it lands, freshness checks detect an index that lags the table. Exact mode defaults to FAIL; approximate mode defaults to WARN and may return stale results with an explicit warning. Callers requiring fail-closed behavior should select `vector.freshness.policy=FAIL`. ########## hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/HoodieVectorSearchPlanBuilder.scala: ########## @@ -17,365 +17,473 @@ package org.apache.spark.sql.hudi.analysis +import org.apache.hudi.SparkAdapterSupport +import org.apache.hudi.client.common.HoodieSparkEngineContext +import org.apache.hudi.common.config.HoodieMetadataConfig +import org.apache.hudi.common.config.TypedProperties +import org.apache.hudi.common.index.vector.{ScoredVectorPostingMatch, VectorDistanceMetric, VectorIndexArbiter, VectorIndexMdtSearchUtils, VectorIndexMetadataCache, VectorIndexOptions, VectorPostingMatch, VectorStalePolicy} +import org.apache.hudi.common.model.{FileSlice, HoodieIndexDefinition, HoodieRecord} import org.apache.hudi.common.schema.HoodieSchema - -import org.apache.spark.sql.{AnalysisException, DataFrame, SparkSession} +import org.apache.hudi.common.table.{HoodieTableMetaClient, TableSchemaResolver} +import org.apache.hudi.common.table.timeline.{HoodieTimeline, InstantComparison} +import org.apache.hudi.common.table.timeline.InstantComparison.compareTimestamps +import org.apache.hudi.common.table.view.HoodieTableFileSystemView +import org.apache.hudi.data.HoodieJavaRDD +import org.apache.hudi.hadoop.fs.HadoopFSUtils +import org.apache.hudi.metadata.{HoodieTableMetadata, HoodieTableMetadataUtil} + +import org.apache.spark.api.java.JavaSparkContext +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SparkSession} +import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.parser.ParseException +import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction.{DistanceMetric, SearchAlgorithm} import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions.{broadcast, col, monotonically_increasing_id, row_number} import org.apache.spark.sql.hudi.command.exception.HoodieAnalysisException -import org.apache.spark.sql.types.{ArrayType, ByteType, DataType, DoubleType, FloatType} +import org.apache.spark.sql.types.{ArrayType, ByteType, DataType, DoubleType, FloatType, Metadata, StructField, StructType} +import org.apache.spark.util.LongAccumulator +import org.slf4j.LoggerFactory +import scala.collection.JavaConverters._ +import scala.collection.mutable.ArrayBuffer import scala.util.{Failure, Success, Try} /** - * Extension point for vector search algorithms. Each implementation provides - * the Spark logical plan for single-query and batch-query KNN search. - * - * To add a new algorithm (e.g. RowMatrix, HNSW): - * 1. Create an object extending this trait - * 2. Add a value to [[SearchAlgorithm]] - * 3. Register the mapping in [[HoodieVectorSearchPlanBuilder.resolveAlgorithm]] + * IVF + RaBitQ vector search backed by native metadata-table posting prefix lookups. * - * Implementations can use the shared validation helpers on - * [[HoodieVectorSearchPlanBuilder]] (validateEmbeddingColumn, validateBatchDimensions, etc.) - * and the raw distance functions on [[VectorDistanceUtils]]. - * - * The output schema contract: - * - Single-query: all corpus columns (minus the embedding column) + `_hudi_distance: Double` - * - Batch-query: all corpus columns (minus the embedding column) + clashing query columns - * (prefixed with `_hudi_query_`) + `_hudi_distance: Double` + `_hudi_query_index: Long` - * - Results are ordered by `_hudi_distance` ascending (lower = more similar) - * - `_hudi_query_index` is an opaque grouping identifier (not a sequential index). Values may be - * large non-contiguous numbers because they are generated by `monotonically_increasing_id()`. + * The algorithm keeps Spark SQL out of the posting scan: it probes centroids, converts + * the selected clusters to MDT posting prefixes, reads posting rows with HFile-friendly + * prefix lookups, reduces with RaBitQ approximate scores before record-location lookup, + * then exact-reranks the fetched Hudi rows. */ -trait VectorSearchAlgorithm { +object IvfRaBitQMdtSearchAlgorithm extends VectorSearchAlgorithm with SparkAdapterSupport { + private val APPROXIMATE_CANDIDATE_FLOOR = 32 - /** Human-readable name for error messages and logging. */ - def name: String + private[analysis] def approximateCandidateHeapSize(k: Int): Int = + math.max(k, APPROXIMATE_CANDIDATE_FLOOR) - /** - * Build a plan that finds the k nearest corpus rows to a single query vector. - * - * @param spark active SparkSession - * @param corpusDf resolved corpus DataFrame (may be Hudi, Parquet, or temp view) - * @param embeddingCol name of the array-typed embedding column in corpusDf - * @param queryVector the query vector, normalized to Array[Double] - * @param k number of nearest neighbors to return - * @param metric distance metric (COSINE, L2, DOT_PRODUCT) - * @param filter optional SQL predicate applied to the corpus before distance computation, - * String literals inside the predicate must use double quotes (e.g. {@code label = "z-axis"}) - * because the filter is passed as a single-quoted TVF argument and - * Spark's lexer does not support {@code ''} as an escape sequence - * inside string literals. - * @param maxDistance optional distance threshold; results exceeding this value are excluded - * @return an analyzed LogicalPlan whose output matches the single-query schema contract - */ - def buildSingleQueryPlan( - spark: SparkSession, - corpusDf: DataFrame, - embeddingCol: String, - queryVector: Array[Double], - k: Int, - metric: DistanceMetric.Value, - filter: Option[String] = None, - maxDistance: Option[Double] = None): LogicalPlan - - /** - * Build a plan that finds the k nearest corpus rows for each row in the query table. - * - * @param spark active SparkSession - * @param corpusDf resolved corpus DataFrame - * @param corpusEmbeddingCol name of the embedding column in corpusDf - * @param queryDf resolved query DataFrame - * @param queryEmbeddingCol name of the embedding column in queryDf - * @param k number of nearest neighbors per query - * @param metric distance metric (COSINE, L2, DOT_PRODUCT) - * @param filter optional SQL predicate applied to the corpus before distance - * computation, and shrinking cross-join cardinality. Applied to the - * corpus only — to restrict the query side, apply a projection or - * filter on the query table before invoking the TVF. See - * [[buildSingleQueryPlan]] for quoting requirements. - * @param maxDistance optional distance threshold; results exceeding this value are - * excluded before per-query top-K selection, reducing shuffle volume. - * @return an analyzed LogicalPlan whose output matches the batch-query schema contract - * @note Batch mode broadcasts the query table to all executors via a cross-join. - * This is designed for small-to-medium query sets (tens to low hundreds of rows). - * For large query tables, memory pressure on executors may occur. - */ - def buildBatchQueryPlan( - spark: SparkSession, - corpusDf: DataFrame, - corpusEmbeddingCol: String, - queryDf: DataFrame, - queryEmbeddingCol: String, - k: Int, - metric: DistanceMetric.Value, - filter: Option[String] = None, - maxDistance: Option[Double] = None): LogicalPlan -} + import HoodieVectorSearchPlanBuilder._ -/** - * Resolves [[SearchAlgorithm]] values to [[VectorSearchAlgorithm]] implementations - * and provides shared validation helpers used across algorithms. - */ -object HoodieVectorSearchPlanBuilder { - - val DISTANCE_COL = "_hudi_distance" - private[analysis] val QUERY_ID_COL = "_hudi_query_index" - private[analysis] val QUERY_EMB_ALIAS = "_hudi_query_emb" - private[analysis] val RANK_COL = "_hudi_rank" - private[analysis] val QUERY_COL_PREFIX = "_hudi_query_" - - /** Resolve a [[SearchAlgorithm]] enum value to its implementation. */ - def resolveAlgorithm(algorithm: SearchAlgorithm.Value): VectorSearchAlgorithm = algorithm match { - case SearchAlgorithm.BRUTE_FORCE => BruteForceSearchAlgorithm - case other => throw new HoodieAnalysisException( - s"Unsupported search algorithm: $other") + private[analysis] def isLogResidentCandidate( + candidate: ScoredVectorPostingMatch, + currentBaseInstant: String, + sliceHasLogs: Boolean): Boolean = { + // A negative position is the authoritative log-resident signal: base-file records always carry a + // non-negative row position, so a candidate with position < 0 in a slice that still has log files + // lives in a log block and cannot be materialized from the base Parquet. The candidate's instant + // is NOT a reliable discriminator here -- a log delta can share the slice's base instant time. + val location = candidate.getLocation + sliceHasLogs && location != null && location.getPosition < 0 } - private[analysis] def validateEmbeddingColumn(df: DataFrame, colName: String): Unit = { - val fieldOpt = df.schema.fields.find(_.name.equalsIgnoreCase(colName)) - val field = fieldOpt.getOrElse( - throw new HoodieAnalysisException( - s"Embedding column '$colName' not found in table schema. " + - s"Available columns: ${df.schema.fieldNames.mkString(", ")}")) - field.dataType match { - case ArrayType(FloatType, _) | ArrayType(DoubleType, _) | ArrayType(ByteType, _) => // valid - case other => - throw new HoodieAnalysisException( - s"Embedding column '$colName' has type $other, " + - "expected array<float>, array<double>, or array<byte>") + private[analysis] def approximateArbitratedCandidate( + candidate: ScoredVectorPostingMatch): Option[ScoredVectorPostingMatch] = + candidate.getArbiterDecision match { + case VectorIndexArbiter.Decision.SERVE => Some(candidate) + // A delta carries the updated vector code. Refreshing only its locator is safe and preserves + // the new score; a stale packed-block row carries old code and must never be resurrected. + case VectorIndexArbiter.Decision.STALE if candidate.isDelta => Some(candidate) + case _ => None } + + private[analysis] case class FreshnessLag(firstUnmarkedInstant: String, lagCount: Int) + + private[analysis] def completedSourceWriteTimeline(timeline: HoodieTimeline): HoodieTimeline = + timeline.getWriteTimeline.filterCompletedInstants() + + private[analysis] def freshnessLag( + completedWriteInstants: Seq[String], + coveredInstant: Option[String]): Option[FreshnessLag] = { + val uncovered = completedWriteInstants.filter(instant => coveredInstant.forall(covered => Review Comment: **Fixed (bootstrap)** `41bdd71c5b`; **follow-up (maintenance)** https://github.com/chrevanthreddy/hudi/pull/5. Bootstrap baseline uses the contiguous completed-write cutoff (`getContiguousCompletedWriteTimeline`), consistent with async index scheduling. Base excludes in-flight writes from bootstrap coverage. Source-instant marker maintenance and advancement through the contiguous covered run are in the stacked update PR, not this bootstrap-only PR. ########## hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/HoodieVectorSearchPlanBuilder.scala: ########## @@ -17,365 +17,473 @@ package org.apache.spark.sql.hudi.analysis +import org.apache.hudi.SparkAdapterSupport +import org.apache.hudi.client.common.HoodieSparkEngineContext +import org.apache.hudi.common.config.HoodieMetadataConfig +import org.apache.hudi.common.config.TypedProperties +import org.apache.hudi.common.index.vector.{ScoredVectorPostingMatch, VectorDistanceMetric, VectorIndexArbiter, VectorIndexMdtSearchUtils, VectorIndexMetadataCache, VectorIndexOptions, VectorPostingMatch, VectorStalePolicy} +import org.apache.hudi.common.model.{FileSlice, HoodieIndexDefinition, HoodieRecord} import org.apache.hudi.common.schema.HoodieSchema - -import org.apache.spark.sql.{AnalysisException, DataFrame, SparkSession} +import org.apache.hudi.common.table.{HoodieTableMetaClient, TableSchemaResolver} +import org.apache.hudi.common.table.timeline.{HoodieTimeline, InstantComparison} +import org.apache.hudi.common.table.timeline.InstantComparison.compareTimestamps +import org.apache.hudi.common.table.view.HoodieTableFileSystemView +import org.apache.hudi.data.HoodieJavaRDD +import org.apache.hudi.hadoop.fs.HadoopFSUtils +import org.apache.hudi.metadata.{HoodieTableMetadata, HoodieTableMetadataUtil} + +import org.apache.spark.api.java.JavaSparkContext +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.{AnalysisException, DataFrame, Row, SparkSession} +import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.parser.ParseException +import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction.{DistanceMetric, SearchAlgorithm} import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.expressions.Window import org.apache.spark.sql.functions.{broadcast, col, monotonically_increasing_id, row_number} import org.apache.spark.sql.hudi.command.exception.HoodieAnalysisException -import org.apache.spark.sql.types.{ArrayType, ByteType, DataType, DoubleType, FloatType} +import org.apache.spark.sql.types.{ArrayType, ByteType, DataType, DoubleType, FloatType, Metadata, StructField, StructType} +import org.apache.spark.util.LongAccumulator +import org.slf4j.LoggerFactory +import scala.collection.JavaConverters._ +import scala.collection.mutable.ArrayBuffer import scala.util.{Failure, Success, Try} /** - * Extension point for vector search algorithms. Each implementation provides - * the Spark logical plan for single-query and batch-query KNN search. - * - * To add a new algorithm (e.g. RowMatrix, HNSW): - * 1. Create an object extending this trait - * 2. Add a value to [[SearchAlgorithm]] - * 3. Register the mapping in [[HoodieVectorSearchPlanBuilder.resolveAlgorithm]] + * IVF + RaBitQ vector search backed by native metadata-table posting prefix lookups. * - * Implementations can use the shared validation helpers on - * [[HoodieVectorSearchPlanBuilder]] (validateEmbeddingColumn, validateBatchDimensions, etc.) - * and the raw distance functions on [[VectorDistanceUtils]]. - * - * The output schema contract: - * - Single-query: all corpus columns (minus the embedding column) + `_hudi_distance: Double` - * - Batch-query: all corpus columns (minus the embedding column) + clashing query columns - * (prefixed with `_hudi_query_`) + `_hudi_distance: Double` + `_hudi_query_index: Long` - * - Results are ordered by `_hudi_distance` ascending (lower = more similar) - * - `_hudi_query_index` is an opaque grouping identifier (not a sequential index). Values may be - * large non-contiguous numbers because they are generated by `monotonically_increasing_id()`. + * The algorithm keeps Spark SQL out of the posting scan: it probes centroids, converts + * the selected clusters to MDT posting prefixes, reads posting rows with HFile-friendly + * prefix lookups, reduces with RaBitQ approximate scores before record-location lookup, + * then exact-reranks the fetched Hudi rows. */ -trait VectorSearchAlgorithm { +object IvfRaBitQMdtSearchAlgorithm extends VectorSearchAlgorithm with SparkAdapterSupport { + private val APPROXIMATE_CANDIDATE_FLOOR = 32 - /** Human-readable name for error messages and logging. */ - def name: String + private[analysis] def approximateCandidateHeapSize(k: Int): Int = + math.max(k, APPROXIMATE_CANDIDATE_FLOOR) - /** - * Build a plan that finds the k nearest corpus rows to a single query vector. - * - * @param spark active SparkSession - * @param corpusDf resolved corpus DataFrame (may be Hudi, Parquet, or temp view) - * @param embeddingCol name of the array-typed embedding column in corpusDf - * @param queryVector the query vector, normalized to Array[Double] - * @param k number of nearest neighbors to return - * @param metric distance metric (COSINE, L2, DOT_PRODUCT) - * @param filter optional SQL predicate applied to the corpus before distance computation, - * String literals inside the predicate must use double quotes (e.g. {@code label = "z-axis"}) - * because the filter is passed as a single-quoted TVF argument and - * Spark's lexer does not support {@code ''} as an escape sequence - * inside string literals. - * @param maxDistance optional distance threshold; results exceeding this value are excluded - * @return an analyzed LogicalPlan whose output matches the single-query schema contract - */ - def buildSingleQueryPlan( - spark: SparkSession, - corpusDf: DataFrame, - embeddingCol: String, - queryVector: Array[Double], - k: Int, - metric: DistanceMetric.Value, - filter: Option[String] = None, - maxDistance: Option[Double] = None): LogicalPlan - - /** - * Build a plan that finds the k nearest corpus rows for each row in the query table. - * - * @param spark active SparkSession - * @param corpusDf resolved corpus DataFrame - * @param corpusEmbeddingCol name of the embedding column in corpusDf - * @param queryDf resolved query DataFrame - * @param queryEmbeddingCol name of the embedding column in queryDf - * @param k number of nearest neighbors per query - * @param metric distance metric (COSINE, L2, DOT_PRODUCT) - * @param filter optional SQL predicate applied to the corpus before distance - * computation, and shrinking cross-join cardinality. Applied to the - * corpus only — to restrict the query side, apply a projection or - * filter on the query table before invoking the TVF. See - * [[buildSingleQueryPlan]] for quoting requirements. - * @param maxDistance optional distance threshold; results exceeding this value are - * excluded before per-query top-K selection, reducing shuffle volume. - * @return an analyzed LogicalPlan whose output matches the batch-query schema contract - * @note Batch mode broadcasts the query table to all executors via a cross-join. - * This is designed for small-to-medium query sets (tens to low hundreds of rows). - * For large query tables, memory pressure on executors may occur. - */ - def buildBatchQueryPlan( - spark: SparkSession, - corpusDf: DataFrame, - corpusEmbeddingCol: String, - queryDf: DataFrame, - queryEmbeddingCol: String, - k: Int, - metric: DistanceMetric.Value, - filter: Option[String] = None, - maxDistance: Option[Double] = None): LogicalPlan -} + import HoodieVectorSearchPlanBuilder._ -/** - * Resolves [[SearchAlgorithm]] values to [[VectorSearchAlgorithm]] implementations - * and provides shared validation helpers used across algorithms. - */ -object HoodieVectorSearchPlanBuilder { - - val DISTANCE_COL = "_hudi_distance" - private[analysis] val QUERY_ID_COL = "_hudi_query_index" - private[analysis] val QUERY_EMB_ALIAS = "_hudi_query_emb" - private[analysis] val RANK_COL = "_hudi_rank" - private[analysis] val QUERY_COL_PREFIX = "_hudi_query_" - - /** Resolve a [[SearchAlgorithm]] enum value to its implementation. */ - def resolveAlgorithm(algorithm: SearchAlgorithm.Value): VectorSearchAlgorithm = algorithm match { - case SearchAlgorithm.BRUTE_FORCE => BruteForceSearchAlgorithm - case other => throw new HoodieAnalysisException( - s"Unsupported search algorithm: $other") + private[analysis] def isLogResidentCandidate( + candidate: ScoredVectorPostingMatch, + currentBaseInstant: String, + sliceHasLogs: Boolean): Boolean = { + // A negative position is the authoritative log-resident signal: base-file records always carry a + // non-negative row position, so a candidate with position < 0 in a slice that still has log files + // lives in a log block and cannot be materialized from the base Parquet. The candidate's instant + // is NOT a reliable discriminator here -- a log delta can share the slice's base instant time. + val location = candidate.getLocation + sliceHasLogs && location != null && location.getPosition < 0 } - private[analysis] def validateEmbeddingColumn(df: DataFrame, colName: String): Unit = { - val fieldOpt = df.schema.fields.find(_.name.equalsIgnoreCase(colName)) - val field = fieldOpt.getOrElse( - throw new HoodieAnalysisException( - s"Embedding column '$colName' not found in table schema. " + - s"Available columns: ${df.schema.fieldNames.mkString(", ")}")) - field.dataType match { - case ArrayType(FloatType, _) | ArrayType(DoubleType, _) | ArrayType(ByteType, _) => // valid - case other => - throw new HoodieAnalysisException( - s"Embedding column '$colName' has type $other, " + - "expected array<float>, array<double>, or array<byte>") + private[analysis] def approximateArbitratedCandidate( + candidate: ScoredVectorPostingMatch): Option[ScoredVectorPostingMatch] = + candidate.getArbiterDecision match { + case VectorIndexArbiter.Decision.SERVE => Some(candidate) + // A delta carries the updated vector code. Refreshing only its locator is safe and preserves + // the new score; a stale packed-block row carries old code and must never be resurrected. + case VectorIndexArbiter.Decision.STALE if candidate.isDelta => Some(candidate) + case _ => None } + + private[analysis] case class FreshnessLag(firstUnmarkedInstant: String, lagCount: Int) + + private[analysis] def completedSourceWriteTimeline(timeline: HoodieTimeline): HoodieTimeline = + timeline.getWriteTimeline.filterCompletedInstants() + + private[analysis] def freshnessLag( + completedWriteInstants: Seq[String], + coveredInstant: Option[String]): Option[FreshnessLag] = { + val uncovered = completedWriteInstants.filter(instant => coveredInstant.forall(covered => + compareTimestamps(covered, InstantComparison.LESSER_THAN, instant))) + uncovered.headOption.map(first => FreshnessLag(first, uncovered.size)) } - /** - * Validates that the query vector dimension matches the corpus embedding dimension - * when the corpus column has VECTOR(dim) metadata. - */ - private[analysis] def validateQueryVectorDimension( - df: DataFrame, embeddingCol: String, queryDim: Int): Unit = { - extractVectorDimension(df, embeddingCol).foreach { corpusDim => - if (corpusDim != queryDim) { - throw new HoodieAnalysisException( - s"Query vector dimension ($queryDim) does not match " + - s"corpus embedding dimension ($corpusDim) for column '$embeddingCol'") + override val name: String = "ivf_rabitq_mdt" + + private[analysis] case class VectorSearchPlanContext( + basePath: String, + metaClient: HoodieTableMetaClient, + timeline: HoodieTimeline, + latestInstantTime: String, + indexPartition: String, + indexDefinition: HoodieIndexDefinition, + metadataTable: HoodieTableMetadata, + cache: VectorIndexMetadataCache, + tuningOptions: Map[String, String], + metric: DistanceMetric.Value) + + private[analysis] def isPartitionedRecordIndex(ctx: VectorSearchPlanContext): Boolean = + false + + private val STRUCTURAL_OPTION_KEYS: Set[String] = Set( + VectorIndexOptions.METRIC, + VectorIndexOptions.QUANTIZER, + VectorIndexOptions.RABITQ_BITS, + VectorIndexOptions.RABITQ_SEED, + VectorIndexOptions.RABITQ_ASSUME_NORMALIZED) + + private[analysis] def resolvedTuningOptions( + indexOptions: Map[String, String], + runtimeOptions: Map[String, String]): Map[String, String] = { + val inherited = indexOptions + .filterNot { case (key, _) => STRUCTURAL_OPTION_KEYS.contains(key) } + // An explicit runtime mode change gets that mode's freshness default unless the caller also + // overrides freshness. Otherwise a persisted exact-mode default (FAIL) would leak into an + // approximate query whose contract default is WARN. + .filterNot { case (key, _) => + key == VectorIndexOptions.FRESHNESS_POLICY && + runtimeOptions.contains(VectorIndexOptions.QUERY_MODE) && + !runtimeOptions.contains(VectorIndexOptions.FRESHNESS_POLICY) } - } + inherited ++ runtimeOptions } - /** - * Validates that corpus and query embedding columns have the same element type. - * - * This is a deliberate design choice: we require strict type matching rather than - * silently upcasting (e.g. float to double). While some vector search systems - * auto-convert, strict matching avoids subtle precision differences and makes - * type mismatches between pipelines explicit. This can be relaxed to support - * automatic widening (byte -> float -> double) in a future iteration if needed. - */ - private[analysis] def validateElementTypeCompatibility( - corpusDf: DataFrame, corpusCol: String, - queryDf: DataFrame, queryCol: String): Unit = { - val corpusElemType = getElementType(corpusDf, corpusCol) - val queryElemType = getElementType(queryDf, queryCol) - if (corpusElemType != queryElemType) { + private def resolvePlanContext( + spark: SparkSession, + basePath: String, + embeddingCol: String, + vectorSchema: HoodieSchema.Vector, + callerMetric: DistanceMetric.Value, + runtimeOptions: Map[String, String]): VectorSearchPlanContext = { + val illegal = runtimeOptions.keySet.intersect(STRUCTURAL_OPTION_KEYS) + if (illegal.nonEmpty) { throw new HoodieAnalysisException( - s"Corpus embedding column '$corpusCol' has element type $corpusElemType " + - s"but query embedding column '$queryCol' has element type $queryElemType. " + - "Both must use the same element type (e.g. array<float>).") + s"Options ${illegal.toSeq.sorted.mkString(", ")} describe vector index structure and cannot be set at query time; " + + "they are read from the active index generation.") } - } - /** - * Validates that corpus and query embedding dimensions match when both have - * VECTOR(dim) metadata. - */ - private[analysis] def validateBatchDimensions( - corpusDf: DataFrame, corpusCol: String, - queryDf: DataFrame, queryCol: String): Unit = { - (extractVectorDimension(corpusDf, corpusCol), extractVectorDimension(queryDf, queryCol)) match { - case (Some(corpusDim), Some(queryDim)) if corpusDim != queryDim => - throw new HoodieAnalysisException( - s"Corpus embedding dimension ($corpusDim) does not match " + - s"query embedding dimension ($queryDim)") - case _ => // dimensions match or metadata not available + val storageConf = HadoopFSUtils.getStorageConf(spark.sessionState.newHadoopConf()) + val metaClient = HoodieTableMetaClient.builder() + .setConf(storageConf.newInstance()) + .setBasePath(basePath) + .build() + // Freshness markers are emitted for completed source writes only. Keep the query comparand on + // exactly that action family: clean, rollback, savepoint, restore, and indexing instants must + // never make an otherwise-covered vector generation look stale. + val timeline = completedSourceWriteTimeline(metaClient.getActiveTimeline) + val latestInstantTime = timeline.lastInstant().orElseThrow(() => + new HoodieAnalysisException(s"No completed commits found for table '$basePath'.")).requestedTime() + val (indexPartition, indexDefinition) = resolveVectorIndexDefinition(metaClient, embeddingCol) + // Build the metadata-reader config from the active session properties instead of bare + // defaults. Prior to this, newBuilder().enable(true).build() silently dropped every + // hoodie.* metadata/hfile read-side conf set on the session (e.g. the HFile block-cache + // sizing), so the vector TVF always read on defaults. fromProperties restores operator + // control; enable(true) still forces MDT on for the search path regardless of session. + val metadataProps = TypedProperties.fromMap(spark.sessionState.conf.getAllConfs.asJava) + val metadataConfig = HoodieMetadataConfig.newBuilder() + .fromProperties(metadataProps) + .enable(true) + .build() + val engineContext = new HoodieSparkEngineContext(new JavaSparkContext(spark.sparkContext), spark.sqlContext) + val metadataTable = metaClient.getTableFormat.getMetadataFactory + .create(engineContext, metaClient.getStorage, metadataConfig, basePath) + val cache = getOrLoadMetadataCache(metaClient, metadataTable, indexPartition, vectorSchema, latestInstantTime) + if (cache == null || cache.getGenerationId < 0) { + throw new HoodieAnalysisException( + s"Vector index '$indexPartition' does not contain valid generation metadata.") } + + val indexOptions = Option(indexDefinition.getIndexOptions).map(_.asScala.toMap).getOrElse(Map.empty) + val indexMetric = toSqlMetric(VectorIndexOptions.getMetric(indexOptions.asJava)) + val authoritativeMetric = validateMetricAuthority(callerMetric, indexMetric, cache) + val tuningOptions = resolvedTuningOptions(indexOptions, runtimeOptions) + VectorSearchPlanContext( + basePath, metaClient, timeline, latestInstantTime, indexPartition, indexDefinition, + metadataTable, cache, tuningOptions, authoritativeMetric) } - private[analysis] def getElementType(df: DataFrame, colName: String): DataType = { - val field = df.schema.fields.find(_.name.equalsIgnoreCase(colName)).getOrElse( - throw new HoodieAnalysisException( - s"Embedding column '$colName' not found in table schema. " + - s"Available columns: ${df.schema.fieldNames.mkString(", ")}")) - field.dataType match { - case ArrayType(elemType, _) => elemType - case other => + private def validateMetricAuthority( + callerMetric: DistanceMetric.Value, + indexMetric: DistanceMetric.Value, + cache: VectorIndexMetadataCache): DistanceMetric.Value = { + if (callerMetric == indexMetric) { + callerMetric + } else { + val normalizedEquivalent = cache.isAssumeNormalized && + Set(callerMetric, indexMetric) == Set(DistanceMetric.COSINE, DistanceMetric.DOT_PRODUCT) + if (normalizedEquivalent) { + callerMetric + } else { throw new HoodieAnalysisException( - s"Embedding column '$colName' has type $other, expected an array type") + s"Query requests metric '$callerMetric' but the vector index generation was built for " + + s"'$indexMetric' (assumeNormalized=${cache.isAssumeNormalized}).") + } } } - /** Extracts VECTOR(dim) dimension from column metadata, if present. */ - private def extractVectorDimension(df: DataFrame, colName: String): Option[Int] = { - df.schema.fields.find(_.name.equalsIgnoreCase(colName)).flatMap { field => - val meta = field.metadata - if (meta.contains(HoodieSchema.TYPE_METADATA_FIELD)) { - val typeDesc = meta.getString(HoodieSchema.TYPE_METADATA_FIELD) - Try(HoodieSchema.parseTypeDescriptor(typeDesc)) match { - case Success(v: HoodieSchema.Vector) => Some(v.getDimension) - case Success(_) => throw new HoodieAnalysisException( - s"Column '$colName' has type '$typeDesc' which is not a VECTOR type. " + - "Only VECTOR columns are supported for vector search.") - case Failure(e) => throw new HoodieAnalysisException( - s"Column '$colName' has malformed type metadata '$typeDesc': ${e.getMessage}") - } - } else None - } + private def toSqlMetric(metric: VectorDistanceMetric): DistanceMetric.Value = metric match { + case VectorDistanceMetric.COSINE => DistanceMetric.COSINE + case VectorDistanceMetric.L2 => DistanceMetric.L2 + case VectorDistanceMetric.DOT_PRODUCT => DistanceMetric.DOT_PRODUCT } -} -/** - * Brute-force KNN vector search: computes distance between every corpus row and - * every query vector, then selects the top-K closest. - * - * <p>Complexity: O(|corpus| * |queries| * dimensions) — linear scan with no index. - * - * <p><b>Single-query mode:</b> applies a distance UDF per corpus row, then - * {@code orderBy + limit(k)} (Spark optimizes this to a partial sort via TakeOrderedAndProject). - * - * <p><b>Batch-query mode:</b> broadcast cross-joins the (small) query table with - * the corpus, computes pairwise distances, then uses a window function to rank - * and select top-K per query. The cross-join produces O(|corpus| * |queries|) - * intermediate rows, so this is suitable for small-to-medium query sets - * (tens to low hundreds of queries) against moderate corpora. - * - * <p>Both modes support an optional {@code filter} predicate (applied to the corpus before - * distance computation), and an optional {@code maxDistance} threshold (results beyond this - * distance are excluded before top-K selection, reducing shuffle and sort volume). - */ -object BruteForceSearchAlgorithm extends VectorSearchAlgorithm { - - import HoodieVectorSearchPlanBuilder._ - - override val name: String = "brute_force" - - /** - * Applies a user-supplied SQL predicate to the corpus DataFrame, wrapping - * [[ParseException]] (predicate syntax error) and [[AnalysisException]] - * (unknown column, type mismatch, etc.) in a [[HoodieAnalysisException]] that - * echoes the offending expression. Other exception types are allowed to - * propagate untouched so they aren't misreported as a filter problem. - */ - private def applyFilter(df: DataFrame, filterExpr: String): DataFrame = { - try { - df.filter(filterExpr) - } catch { - case e @ (_: ParseException | _: AnalysisException) => - throw new HoodieAnalysisException( - s"Invalid pre-filter expression '$filterExpr': ${e.getMessage}") - } + private[analysis] def toVectorMetric(metric: DistanceMetric.Value): VectorDistanceMetric = metric match { + case DistanceMetric.COSINE => VectorDistanceMetric.COSINE + case DistanceMetric.L2 => VectorDistanceMetric.L2 + case DistanceMetric.DOT_PRODUCT => VectorDistanceMetric.DOT_PRODUCT } override def buildSingleQueryPlan( spark: SparkSession, - corpusDf: DataFrame, + corpusTable: VectorSearchTable, embeddingCol: String, queryVector: Array[Double], k: Int, metric: DistanceMetric.Value, - filter: Option[String], - maxDistance: Option[Double]): LogicalPlan = { + runtimeOptions: Map[String, String] = Map.empty): LogicalPlan = { + val corpusDf = corpusTable.df validateEmbeddingColumn(corpusDf, embeddingCol) validateQueryVectorDimension(corpusDf, embeddingCol, queryVector.length) - val elemType = getElementType(corpusDf, embeddingCol) - // Apply pre-filter before distance computation to enable reducing the number of rows that need distance computation. - var filteredDf = corpusDf.filter(col(embeddingCol).isNotNull) - filter.foreach(f => filteredDf = applyFilter(filteredDf, f)) - - // Validate byte corpus query vector values before creating the UDF. - if (elemType == ByteType) { - queryVector.foreach { v => - if (v < Byte.MinValue || v > Byte.MaxValue) - throw new HoodieAnalysisException( - s"Query vector value $v is out of range for byte corpus [${Byte.MinValue}, ${Byte.MaxValue}]") - if (v != v.toByte.toDouble) - throw new HoodieAnalysisException( - s"Query vector value $v is not a whole number. " + - "Byte corpus embeddings require integer-valued query vectors.") - } + val basePath = corpusTable.basePath.getOrElse { + throw new HoodieAnalysisException( + s"Search algorithm '$name' requires a Hudi table identifier or table path, not a temporary view.") + } + if (!corpusDf.columns.contains(HoodieRecord.RECORD_KEY_METADATA_FIELD)) { + throw new HoodieAnalysisException( + s"Search algorithm '$name' requires Hudi meta fields, but '${HoodieRecord.RECORD_KEY_METADATA_FIELD}' is missing.") } - // The single-query UDF closes over the pre-computed query DenseVector, - // so only the corpus column is passed per row. - val distanceUdf = VectorDistanceUtils.createSingleQueryDistanceUdf(metric, elemType, queryVector) - - var scored = filteredDf - .withColumn(DISTANCE_COL, distanceUdf(col(embeddingCol))) - .drop(embeddingCol) - // Apply max-distance threshold before ordering to shrink the sort input. - maxDistance.foreach(d => scored = scored.filter(col(DISTANCE_COL) <= d)) - val result = scored.orderBy(col(DISTANCE_COL).asc).limit(k) - - result.queryExecution.analyzed + val vectorSchema = extractVectorSchema(corpusDf, embeddingCol).getOrElse { + throw new HoodieAnalysisException( + s"Search algorithm '$name' requires '$embeddingCol' to carry VECTOR metadata.") + } + val queryFloat = queryVector.map(_.toFloat) + + val contextStartNs = System.nanoTime() + val ctx = resolvePlanContext(spark, basePath, embeddingCol, vectorSchema, metric, runtimeOptions) + LOG.info( + s"[vector_search][plan_context] basePath=$basePath latestInstant=${ctx.latestInstantTime} " + + s"elapsedMs=${elapsedMs(contextStartNs)}") + val coveredInstant = ctx.cache.getLastContiguousSourceInstant + val lag = freshnessLag( + ctx.timeline.getInstantsAsStream.iterator.asScala.map(_.requestedTime()).toSeq, + Option(coveredInstant)) + val staleFallbackPlan = lag.map { freshnessLag => + val freshnessPolicy = VectorIndexOptions.getFreshnessPolicy(ctx.tuningOptions.asJava) + val firstUnmarkedInstant = freshnessLag.firstUnmarkedInstant + val lagCount = freshnessLag.lagCount + val staleMessage = + s"Vector index generation ${ctx.cache.getGenerationId} covers source writes through " + + s"${Option(coveredInstant).getOrElse("<none>")} but the latest source data-write instant is ${ctx.latestInstantTime}; " + + s"first unmarked instant=$firstUnmarkedInstant, lag count=$lagCount. " + + s"Run vector-index catch-up/replay (or rebuild the index until replay is available) to repair the marker gap" + freshnessPolicy match { + case VectorStalePolicy.FAIL => + throw new HoodieAnalysisException(staleMessage) + case VectorStalePolicy.WARN => + LOG.warn(s"$staleMessage; continuing because vector.freshness.policy=WARN") + None + case VectorStalePolicy.FALLBACK => + LOG.warn(s"$staleMessage; using brute-force exact fallback") + Some(BruteForceSearchAlgorithm.buildSingleQueryPlan( + spark, corpusTable, embeddingCol, queryVector, k, metric, runtimeOptions)) + } + }.flatten + staleFallbackPlan.getOrElse { + if (ctx.cache.getDimension != vectorSchema.getDimension) { + throw new HoodieAnalysisException( + s"Vector index dimension (${ctx.cache.getDimension}) does not match column '$embeddingCol' dimension (${vectorSchema.getDimension}).") + } + val approxMode = VectorIndexOptions.isApproximateSearchMode(ctx.tuningOptions.asJava) + + if (approxMode) { + buildApproximatePlan(spark, ctx, queryFloat, k) + } else { + val exactOutputSchema = StructType( + corpusDf.schema.fields.filterNot(_.name.equalsIgnoreCase(embeddingCol)) :+ + StructField(DISTANCE_COL, DoubleType, nullable = false)) + val candidateDf = VectorExactSearchPlanner.buildRaBitQExactCandidateDf( + spark, ctx, exactOutputSchema, embeddingCol, vectorSchema, queryFloat, queryVector, k) + + candidateDf + .orderBy(col(DISTANCE_COL).asc) + .limit(k) + .queryExecution.analyzed + } + } } override def buildBatchQueryPlan( spark: SparkSession, - corpusDf: DataFrame, + corpusTable: VectorSearchTable, corpusEmbeddingCol: String, queryDf: DataFrame, queryEmbeddingCol: String, k: Int, metric: DistanceMetric.Value, - filter: Option[String], - maxDistance: Option[Double]): LogicalPlan = { - validateEmbeddingColumn(corpusDf, corpusEmbeddingCol) - validateEmbeddingColumn(queryDf, queryEmbeddingCol) - validateElementTypeCompatibility(corpusDf, corpusEmbeddingCol, queryDf, queryEmbeddingCol) - validateBatchDimensions(corpusDf, corpusEmbeddingCol, queryDf, queryEmbeddingCol) - - val corpusElemType = getElementType(corpusDf, corpusEmbeddingCol) - val distanceUdf = VectorDistanceUtils.createDistanceUdf(metric, corpusElemType) - // Apply pre-filter before the cross-join to enable Hudi partition pruning and - // data skipping, reducing the cross-join cardinality. - var filteredCorpus = corpusDf.filter(col(corpusEmbeddingCol).isNotNull) - filter.foreach(f => filteredCorpus = applyFilter(filteredCorpus, f)) - - // Prefix clashing query columns with "_hudi_query_" to avoid cross-join ambiguity when - // corpus and query share column names (e.g. both have "id" or "embedding"). - val corpusCols = filteredCorpus.columns.toSet - val queryWithId = queryDf.filter(col(queryEmbeddingCol).isNotNull) - .withColumnRenamed(queryEmbeddingCol, QUERY_EMB_ALIAS) - .withColumn(QUERY_ID_COL, monotonically_increasing_id()) - - val renamedQuery = queryWithId.select(queryWithId.columns.map { qCol => - if (qCol != QUERY_ID_COL && qCol != QUERY_EMB_ALIAS && corpusCols.contains(qCol)) - col(qCol).as(s"$QUERY_COL_PREFIX$qCol") - else - col(qCol) - }: _*) - - // Cross join corpus with broadcast queries, compute distance, then rank - var scored = filteredCorpus.crossJoin(broadcast(renamedQuery)) - .withColumn(DISTANCE_COL, - distanceUdf(col(corpusEmbeddingCol), col(QUERY_EMB_ALIAS))) - .drop(corpusEmbeddingCol) - .drop(QUERY_EMB_ALIAS) - - // Apply max-distance threshold before windowing to reduce shuffle volume. - maxDistance.foreach(d => scored = scored.filter(col(DISTANCE_COL) <= d)) - - val window = Window.partitionBy(QUERY_ID_COL).orderBy(col(DISTANCE_COL).asc) - val result = scored - .withColumn(RANK_COL, row_number().over(window)) - .filter(col(RANK_COL) <= k) - .drop(RANK_COL) - .orderBy(col(QUERY_ID_COL), col(DISTANCE_COL)) - - result.queryExecution.analyzed + runtimeOptions: Map[String, String] = Map.empty): LogicalPlan = { + throw new HoodieAnalysisException( + s"Search algorithm '$name' currently supports single-query hudi_vector_search only.") + } + + /** + * Approximate-only path: scores postings via RaBitQ and returns + * (record_key, partition_path, file_group_id, approx_distance) without + * reading the base table at all. + */ + private def buildApproximatePlan( + spark: SparkSession, + ctx: VectorSearchPlanContext, + queryVector: Array[Float], + k: Int): LogicalPlan = { + import spark.implicits._ + val planStartNs = System.nanoTime() + val cache = ctx.cache + val mergedOptions = ctx.tuningOptions + val metric = toVectorMetric(ctx.metric) + val residualEncoding = cache.isResidualEncoding + val numProbes = math.min(math.max(1, VectorIndexOptions.getNumProbes(mergedOptions.asJava)), cache.numClusters()) + val centroidProbeStartNs = System.nanoTime() + val topClusters = cache.findTopClusters(queryVector, numProbes, metric) + val centroidProbeMs = elapsedMs(centroidProbeStartNs) + val shardLookupStartNs = System.nanoTime() + val shardCounts = cache.getShardCounts(topClusters) + val shardLookupMs = elapsedMs(shardLookupStartNs) + // Approximate-only search does not exact-rerank, so applying refine_factor here only bloats + // every task-local heap and both global reductions. Keep a small floor because local pruning + // precedes record-key overlay and tombstone removal; retaining exactly k could underfill results + // when duplicate versions occupy the local winners. + val candidateHeapSize = approximateCandidateHeapSize(k) + val postingPlanStartNs = System.nanoTime() + val topCandidates = VectorIndexMdtSearchUtils.scanPostingCandidates( + ctx.metadataTable, ctx.indexPartition, cache.getGenerationId, shardCounts, + queryVector, cache.getDimension, cache.getQuantizerSeed, cache.getRaBitQBits, cache.isAssumeNormalized, + metric, VectorIndexOptions.isRaBitQAsymmetricScoring(mergedOptions.asJava), + residualEncoding, if (residualEncoding) cache.getCentroids else null, candidateHeapSize) + val postingPlanMs = elapsedMs(postingPlanStartNs) + val resultRddPlanStartNs = System.nanoTime() + val resultSchema = new StructType(Array( + StructField(HoodieRecord.RECORD_KEY_METADATA_FIELD, org.apache.spark.sql.types.DataTypes.StringType, nullable = false, Metadata.empty), + StructField(HoodieRecord.PARTITION_PATH_METADATA_FIELD, org.apache.spark.sql.types.DataTypes.StringType, nullable = true, Metadata.empty), + StructField(FILE_GROUP_ID_COL, org.apache.spark.sql.types.DataTypes.StringType, nullable = true, Metadata.empty), + StructField(DISTANCE_COL, org.apache.spark.sql.types.DataTypes.DoubleType, nullable = false, Metadata.empty) + )) + // RFC-109 RLI finalist arbiter (freshness gate). Default-off: dormant until the commit-time Review Comment: **Fixed** `3cde5283ea`. The arbiter is mandatory; the never-disabled flag and its dead branches are removed. Approx mode depends on RLI by design (stale/deleted filtering), and CREATE INDEX requires a record index. -- 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]
