chrevanthreddy commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4384292504
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/HoodieVectorSearchTableValuedFunction.scala:
##########
@@ -91,49 +91,35 @@ object HoodieVectorSearchTableValuedFunction {
else DistanceMetric.COSINE
val algorithm = if (exprs.size >= 6)
SearchAlgorithm.fromString(requireStringLiteral(exprs(5), "algorithm"))
else SearchAlgorithm.BRUTE_FORCE
- val filter = if (exprs.size >= 7) parseOptionalString(FUNC_NAME, exprs(6),
"filter") else None
- val maxDistance = if (exprs.size >= 8) parseOptionalDouble(FUNC_NAME,
exprs(7), "max_distance") else None
- ParsedArgs(table, embeddingCol, queryVectorExpr, k, metric, algorithm,
filter, maxDistance)
- }
-
- private[logical] def parseK(funcName: String, expr: Expression): Int = {
- val rawValue = expr.eval()
- if (rawValue == null) {
- throw new HoodieAnalysisException(
- s"Function '$funcName': k must be a positive integer, got null")
+ val runtimeOptions = algorithm match {
+ case SearchAlgorithm.BRUTE_FORCE =>
+ val filter = if (exprs.size >= 7) parseOptionalString(FUNC_NAME,
exprs(6), "filter") else None
+ val maxDistance = if (exprs.size >= 8) parseOptionalDouble(FUNC_NAME,
exprs(7), "max_distance") else None
+ filter.map(BRUTE_FORCE_FILTER_OPT -> _).toMap ++
+ maxDistance.map(value => BRUTE_FORCE_MAX_DISTANCE_OPT ->
value.toString).toMap
+ case SearchAlgorithm.IVF_RABITQ_MDT =>
+ if (exprs.size >= 8) {
+ throw new HoodieAnalysisException(
+ s"Function '$FUNC_NAME': max_distance is supported only by
brute_force")
+ }
+ if (exprs.size >= 7) parseKvOptions(requireStringLiteral(exprs(6),
"options"))
Review Comment:
**Fixed** `3cde5283ea`. NULL is accepted for the `ivf_rabitq_mdt` options
argument in both TVFs (`TestVectorSearchTvfArguments`).
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/HoodieSparkBaseAnalysis.scala:
##########
@@ -345,12 +343,19 @@ case class ResolveReferences(spark: SparkSession) extends
Rule[LogicalPlan]
* Resolves a table reference to a DataFrame. Accepts either a table
identifier
* (including multi-part identifiers like catalog.db.table) or a file path.
*/
- private def resolveTableToDf(table: String, funcName: String): DataFrame = {
+ private def resolveTableToDf(table: String, funcName: String):
VectorSearchTable = {
try {
if (table.contains(StoragePath.SEPARATOR)) {
- spark.read.format("hudi").load(table)
+ VectorSearchTable(spark.read.format("hudi").load(table), Some(table))
} else {
- spark.table(table)
+ val tableDf = spark.table(table)
+ val basePathOpt = try {
+ val tableId =
spark.sessionState.sqlParser.parseTableIdentifier(table)
Review Comment:
**Fixed** `3cde5283ea`. The index base path is resolved from the resolved
Hudi catalog table (`SparkAdapter.resolveHoodieTable`); a catalog-qualified
test was added.
##########
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
+ // delta writer lands. When enabled, exclude STALE (rewritten/moved) and
DELETED finalists so
+ // approx candidate-gen never serves a posting that no longer faithfully
represents a live row.
+ val arbiterEnabled =
VectorIndexOptions.isFinalistArbiterEnabled(mergedOptions.asJava)
+ val candidateRdd: RDD[ScoredVectorPostingMatch] =
+ if (arbiterEnabled) {
+ val arbitrated = VectorIndexMdtSearchUtils.arbitrateFinalists(
+ ctx.metadataTable, topCandidates, isPartitionedRecordIndex(ctx))
+ val staleAcc: LongAccumulator =
spark.sparkContext.longAccumulator("vector_arbiter_stale")
+ val deletedAcc: LongAccumulator =
spark.sparkContext.longAccumulator("vector_arbiter_deleted")
+ HoodieJavaRDD.getJavaRDD(arbitrated).rdd.mapPartitions { it =>
+ var stale = 0L
+ var deleted = 0L
+ val kept = ArrayBuffer.empty[ScoredVectorPostingMatch]
+ it.foreach { c =>
+ approximateArbitratedCandidate(c) match {
+ case Some(candidate) => kept += candidate
+ case None if c.getArbiterDecision ==
VectorIndexArbiter.Decision.STALE => stale += 1L
+ case None => deleted += 1L
+ }
+ }
+ staleAcc.add(stale)
+ deletedAcc.add(deleted)
+ if (stale > 0L || deleted > 0L) {
+ LOG.info(s"[vector_search][approximate][arbiter]
arbiterExclusions{stale=$stale, deleted=$deleted} kept=${kept.size}")
+ }
+ kept.iterator
+ }
+ } else {
+ HoodieJavaRDD.getJavaRDD(topCandidates).rdd
+ }
+ val resultRows = candidateRdd.map { c =>
+ val liveLocation = Option(c.getLocation)
+ Row(
+ c.getRecordKey,
+
liveLocation.map(_.getPartitionPath).orElse(Option(c.getPartitionPath)).getOrElse(""),
+
liveLocation.map(_.getFileId).orElse(Option(c.getFileGroupId)).getOrElse(""),
+ c.getApproxDistance.toDouble
+ )
+ }
+ val resultRddPlanMs = elapsedMs(resultRddPlanStartNs)
+ val sparkAnalysisStartNs = System.nanoTime()
+ val analyzedPlan = spark.createDataFrame(resultRows, resultSchema)
+ .orderBy(col(DISTANCE_COL).asc)
+ .limit(k)
+ .queryExecution.analyzed
+ val sparkAnalysisMs = elapsedMs(sparkAnalysisStartNs)
+ LOG.info(
+ s"[vector_search][approximate] basePath=${ctx.basePath}
probedClusters=${topClusters.length} " +
+ s"candidateHeapSize=$candidateHeapSize
residualEncoding=$residualEncoding arbiter=$arbiterEnabled " +
+ s"centroidProbeMs=$centroidProbeMs shardLookupMs=$shardLookupMs
postingPlanMs=$postingPlanMs " +
+ s"resultRddPlanMs=$resultRddPlanMs sparkAnalysisMs=$sparkAnalysisMs
planMs=${elapsedMs(planStartNs)}")
+ analyzedPlan
+ }
+
+ private def resolveVectorIndexDefinition(
+ metaClient: HoodieTableMetaClient,
+ embeddingCol: String): (String, HoodieIndexDefinition) = {
+ if (!metaClient.getIndexMetadata.isPresent) {
+ throw new HoodieAnalysisException("No Hudi index metadata is available
for the table.")
+ }
+ val matches = metaClient.getIndexMetadata.get.getIndexDefinitions.asScala
+ .filter { case (_, definition) =>
+ definition.getIndexType ==
HoodieTableMetadataUtil.PARTITION_NAME_VECTOR_INDEX
+ }
+ .filter { case (_, definition) =>
+
Option(definition.getSourceFields).exists(_.asScala.exists(_.equalsIgnoreCase(embeddingCol)))
+ }
+ .toSeq
+ .sortBy(_._1)
+ matches match {
+ case Seq() =>
+ throw new HoodieAnalysisException(
+ s"No vector index found for embedding column '$embeddingCol'.")
+ case Seq(single) =>
+ single
+ case multiple =>
+ throw new HoodieAnalysisException(
+ s"Multiple vector indexes found for embedding column
'$embeddingCol': ${multiple.map(_._1).mkString(", ")}. " +
+ "Specify the index explicitly or drop the stale index.")
+ }
+ }
+
+ private case class MetadataCacheKey(basePath: String, indexPartition: String)
+
+ @volatile private var metadataCaches: Map[MetadataCacheKey,
VectorIndexMetadataCache] = Map.empty
+
+ private[analysis] def resetMetadataCaches(): Unit = {
+ metadataCaches = Map.empty
}
+ private[analysis] def metadataCacheSize: Int = metadataCaches.size
+
+ private def getOrLoadMetadataCache(
+ metaClient: HoodieTableMetaClient,
+ metadataTable: HoodieTableMetadata,
+ indexPartition: String,
+ vectorSchema: HoodieSchema.Vector,
+ currentInstant: String): VectorIndexMetadataCache = {
+ val cacheKey = MetadataCacheKey(metaClient.getBasePath.toString,
indexPartition)
+
+ metadataCaches.synchronized {
Review Comment:
**Fixed** `3cde5283ea`. Loads/replacement go through a dedicated
`metadataCachesLock`; the map is immutable and published via `@volatile`.
##########
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)
Review Comment:
**Fixed** `3cde5283ea`. Unknown query-time keys are rejected; the allow-list
is shared with CREATE INDEX validation (`VectorIndexOptions.QUERY_OPTIONS`).
--
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]
