chrevanthreddy commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4384291881
##########
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)** in `41bdd71c5b`; maintenance follow-up:
https://github.com/chrevanthreddy/hudi/pull/5. Bootstrap uses the contiguous
completed-write cutoff (`getContiguousCompletedWriteTimeline`), consistent with
async index scheduling, rather than treating an inflight write as covered. The
inflight-writer regression verifies the baseline excludes writer A, A emits its
marker when it commits, and exact search includes both A and B.
Update (2026-10-27): the maintenance follow-up no longer advances a
per-write manifest frontier or maintains coverage barriers. Normal writes
maintain postings and an audit marker together; strict read validation is
opt-in/default off, with WARN the default policy in both modes. This does not
remove the existing async-index bootstrap/catch-up lifecycle.
Fixture ordering, resequenced in the maintenance follow-up: observe inflight
partition marking in table config, wait for bootstrap-built state and the
rewrite backup to disappear, then release A. Catch-up keeps final index
completion behind A's commit. The test remains enabled with its original
coverage assertions; all six production E2Es passed with retries disabled on
Spark 3.5 and Spark 4.0 at `4097189466`. No config-reader production changes
are included.
--
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]
