You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@hudi.apache.org by "chrevanthreddy (via GitHub)" <gi...@apache.org> on 2026/09/27 21:58:47 UTC
[PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
chrevanthreddy opened a new pull request, #20107:
URL: https://github.com/apache/hudi/pull/20107
## Summary
Port the proven RFC-109 IVF/RaBitQ vector-index bootstrap and Spark search implementation onto a clean current-master branch, then mechanically separate the posting scanner, RLI arbitration, and Spark read-path components. This draft is a **review/benchmark staging PR**, not a claim of completed 10M qualification.
Related: RFC-109 #19309; prior incremental draft #19802 (historical implementation; please review this refactor against it rather than treating both as separate features). Partitioned-RLI integration/coverage is tracked separately in #20106.
## Included
- Vector MDT Avro/payload, RaBitQ factor contract, bootstrap posting writer, generation manifest/cache, Spark SQL CREATE INDEX wiring, and option validation.
- Historical packed posting scanner, delta overlay/tombstones, RaBitQ scoring, finalist RLI arbitration, and Spark approximate/exact search with positional Parquet/key-fallback and log-resident fetch paths.
- Brute-force TVF filter/max-distance backward compatibility, with IVF runtime tuning options.
- Mechanical read-path extractions into `VectorSearchAlgorithms`, `VectorExactSearchPlanner`, `VectorIndexRliArbitrator`, posting serde and candidate helpers. The preserved read-path was ported first, then split.
## Scope boundaries / review warnings
- **Not yet 10M benchmarked on this branch.** Historical BigANN SparkApplication YAML targets Spark 4.0 / Scala 2.13 and shared writable GCS paths; this branch builds Spark 3.5 / Scala 2.12. Do not apply old manifests unchanged or attribute historical benchmark results to this branch. Current workstation GCS access requires interactive reauthentication; 10M prerequisites and isolated outputs remain unverified.
- Partitioned RLI detection is intentionally hardcoded `false` in `IvfRaBitQMdtSearchAlgorithm.isPartitionedRecordIndex` for this initial baseline. The existing partition-aware arbiter path exists, but must not be claimed correct for partitioned-RLI tables until #20106 is addressed.
- HFile cache telemetry/performance experiment, range prefetch, and unrelated generic core fixes are excluded. Existing defaults and SI/RLI/bootstrap behavior should not be changed by optional caching work.
- This is 18 commits / 87 files from current merge base because it includes schema/DDL/bootstrap foundations as well as read path. Earlier RFC-109 PRs overlap; reviewers may prefer splitting/rebasing after benchmark validation.
## Validation done (Spark 3.5 / Scala 2.12)
- Reactor `test-compile` passed (16 modules; RAT, Checkstyle, Scalastyle).
- Focused Spark JUnit: planner 8/8, Parquet locator 1/1, TVF argument compatibility 5/5 (14/14 total).
- Existing execution-level TVF tests: single- and batch-query combined filter/max-distance and invalid filter handling: 3/3 passed. Maven continued into unrelated ScalaTest suites; it was stopped after selected tests passed, so this is **not** a full test-suite pass.
- `TestVectorIndexOptions`: 10/10; posting/RaBitQ/arbiter focused tests passed in earlier refactor slices.
- Spark 3.5/Scala 2.12 shaded bundle packaged successfully at `packaging/hudi-spark-bundle/target/hudi-spark3.5-bundle_2.12-1.3.0-SNAPSHOT.jar` (local artifact; not deployed).
## Remaining review gates
- Arrange GCS reauthentication and compatible Spark 3.5/Scala 2.12 cluster runtime; stage the bundle and adapted BigANN app under an isolated, immutable run ID; require source `_SUCCESS` markers.
- Run 10M COW bootstrap/index/query with exact and approximate modes, collect latency/recall and inspect Spark logs plus metrics `_SUCCESS` before claiming success. Add real numbers and artifacts here.
- Confirm non-partitioned RLI lifecycle and broaden regression coverage; #20106 separately tracks partitioned RLI.
Please keep this PR **draft** until the 10M gate and reviewer feedback are resolved.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "vinishjail97 (via GitHub)" <gi...@apache.org>.
vinishjail97 commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-5943444310
Hi @chrevanthreddy, thank you for putting this re-port together, and for the initial vector search PRs.
I tested this branch end to end on lakegres with the Hugging Face Wikipedia dataset (`CREATE INDEX ... USING VECTOR`, then `hudi_vector_search`). I ran into a few issues along the way and put the fixes I needed in a small stacked PR on your branch: chrevanthreddy/hudi#3. It has 9 commits, and most of them are a few lines each:
- accept `vector_index` as an index type (`CREATE INDEX` failed without it)
- record the bootstrap instant in the manifest (the freshness check failed every query right after bootstrap)
- clamp the coarse KMeans cells to `vector.num_clusters`
- store extended RaBitQ planes MSB-first, which the scorer expects for bits >= 3
- validate metric/bits at DDL with an `l2` default, and require COW plus a record index
- handle bits=1 when the dimension is not a multiple of 64
- a new end-to-end `TestHoodieVectorIndexSearch`, so CI covers this flow
Please feel free to take, change, or drop any of these. You know this code better than I do. The PR description also lists a few follow-ups that I did not try to fix: incremental updates, partitioned RLI (#20106), chunked centroids, and bits=1 cosine/dot scoring. I would be glad to pair on any of them.
Thanks again!
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4384295539
##########
hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQEncoder.java:
##########
@@ -0,0 +1,588 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.common.index.vector;
+
+import java.io.Serializable;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * RaBitQ (Randomized Binary Quantization) encoder.
+ *
+ * <p>Encodes float vectors into 1-bit-per-dimension binary codes using a random
+ * orthogonal rotation matrix R. The rotation is deterministic given (seed, dimension),
+ * so R is stored as a seed — it never needs retraining.
+ *
+ * <p><b>Encoding (write path):</b>
+ * <ol>
+ * <li>Normalize: v̂ = v / ||v||</li>
+ * <li>Rotate: v_rot = R @ v̂</li>
+ * <li>Binarize: code = pack(sign(v_rot)) — D bits → ceil(D/8) bytes</li>
+ * <li>Scalar: s = ||v|| (stored alongside code; 1.0 if assume_normalized)</li>
+ * </ol>
+ *
+ * <p><b>Query scan (read path):</b>
+ * <ol>
+ * <li>Normalize + rotate the query vector (same R)</li>
+ * <li>Binarize to get q_bin</li>
+ * <li>For each database code b: Hamming(q_bin, b) → estimated cosine</li>
+ * <li>Re-rank top-R candidates with exact distance</li>
+ * </ol>
+ *
+ * <p>Memory: the rotation matrix is D×D floats = 4·D² bytes (~2.3 MB at D=768).
+ * It is built lazily on first use and reused across encode calls.
+ *
+ * <p>Thread-safe after the first call to {@link #encode} or {@link #encodeQuery}
+ * (the lazy init is synchronized).
+ */
+public final class RaBitQEncoder implements Serializable {
+ private static final long serialVersionUID = 1L;
+ private static final Map<RotationKey, float[][]> ROTATION_MATRIX_CACHE = new ConcurrentHashMap<>();
+ /** Neutral-factor thresholds for the current posting-block format (RFC-109 §3). */
+ private static final RaBitQFactorConfig FACTOR_CONFIG = RaBitQFactorConfig.defaults();
+ private final int dimension;
+ private final int bits;
+ private final long seed;
+ private final boolean assumeNormalized;
+
+ /** Rotation matrix, row-major. Populated lazily. */
+ private transient volatile float[][] rotMat;
+ public RaBitQEncoder(int dimension, long seed, boolean assumeNormalized) {
+ this(dimension, 1, seed, assumeNormalized);
+ }
+
+ public RaBitQEncoder(int dimension, int bits, long seed, boolean assumeNormalized) {
+ if (dimension <= 0) {
+ throw new IllegalArgumentException("Dimension must be positive, got: " + dimension);
+ }
+ if (bits <= 0 || bits > 8) {
+ throw new IllegalArgumentException("RaBitQ bits must be in [1, 8], got: " + bits);
+ }
+ this.dimension = dimension;
+ this.bits = bits;
+ this.seed = seed;
+ this.assumeNormalized = assumeNormalized;
+ }
+
+ /** Convenience constructor using default seed. */
+ public RaBitQEncoder(int dimension) {
+ this(dimension, 1, 42L, false);
+ }
+
+ // ---- encoding ----------------------------------------------------------
+
+ public QuantizedVector encode(float[] vector) {
+ validateVector(vector, "vector");
+ float norm = norm(vector);
+ float scalar = assumeNormalized ? 1.0f : norm;
+ float[] normalized = (norm == 0f || assumeNormalized)
+ ? vector
+ : normalize(vector, norm);
+ float[] rotated = rotate(normalized);
+ return new QuantizedVector(binarize(rotated), scalar);
+ }
+
+ /**
+ * Metric-neutral residual encoding used by MDT postings.
+ *
+ * <p>The optional {@code center} is the IVF centroid in the <b>original</b> vector space.
+ * When null, the vector is encoded relative to the origin.
+ */
+ public QuantizedVector encodeResidual(float[] vector, float[] center) {
Review Comment:
**Fixed** `03218b0f6f`. `fromNeutralFactors` names the only difference between bit widths (total rescale).
##########
hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexMdtSearchUtils.java:
##########
@@ -0,0 +1,1310 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.common.index.vector;
+
+import org.apache.hudi.avro.model.HoodieVectorIndexPostingBlock;
+import org.apache.hudi.common.data.HoodieData;
+import org.apache.hudi.common.data.HoodieListData;
+import org.apache.hudi.common.data.HoodiePairData;
+import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.model.HoodieRecordGlobalLocation;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.metadata.HoodieMetadataPayload;
+import org.apache.hudi.metadata.HoodieTableMetadata;
+import org.apache.hudi.metadata.RawKey;
+import org.apache.hudi.metadata.VectorIndexMetadataKey;
+import org.apache.hudi.metadata.VectorPostingPrefixRawKey;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.management.GarbageCollectorMXBean;
+import java.lang.management.ManagementFactory;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+
+/**
+ * Shared helper for MDT-native vector posting lookup and approximate candidate reduction.
+ *
+ * <p>This provides the missing second-stage building block after coarse IVF cluster pruning:
+ * exact HFile lookups of cluster metadata, prefix scans over {@code P|gen|cluster|shard|},
+ * approximate RaBitQ scoring, and direct file-slice targeting from posting payload metadata.
+ */
+public final class VectorIndexMdtSearchUtils {
+
+ private static final Logger LOG = LoggerFactory.getLogger(VectorIndexMdtSearchUtils.class);
+ private static final int TOP_K_REDUCER_KEY = 0;
+ private static final Comparator<ScoredVectorPostingMatch> BEST_FIRST = Comparator
+ .comparingDouble(ScoredVectorPostingMatch::getApproxDistance)
+ .thenComparing(ScoredVectorPostingMatch::getRecordKey);
+
+ private VectorIndexMdtSearchUtils() {
+ }
+
+ public static Map<Integer, Integer> readClusterShardCounts(HoodieTableMetadata metadataTable,
+ String indexPartition,
+ int generationId,
+ Collection<Integer> clusterIds) {
+ if (clusterIds == null || clusterIds.isEmpty()) {
+ return Collections.emptyMap();
+ }
+
+ Map<Integer, Integer> shardCounts = new HashMap<>();
+ for (Integer clusterId : clusterIds) {
+ shardCounts.put(clusterId, 1);
+ }
+ return shardCounts;
+ }
+
+ public static Map<Integer, Set<String>> readClusterFileGroups(HoodieTableMetadata metadataTable,
Review Comment:
**Fixed** `c8b06883b1`. `readClusterShardCounts` / `readClusterFileGroups` had no callers; removed.
##########
hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexMdtSearchUtils.java:
##########
@@ -0,0 +1,1310 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.common.index.vector;
+
+import org.apache.hudi.avro.model.HoodieVectorIndexPostingBlock;
+import org.apache.hudi.common.data.HoodieData;
+import org.apache.hudi.common.data.HoodieListData;
+import org.apache.hudi.common.data.HoodiePairData;
+import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.model.HoodieRecordGlobalLocation;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.metadata.HoodieMetadataPayload;
+import org.apache.hudi.metadata.HoodieTableMetadata;
+import org.apache.hudi.metadata.RawKey;
+import org.apache.hudi.metadata.VectorIndexMetadataKey;
+import org.apache.hudi.metadata.VectorPostingPrefixRawKey;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.management.GarbageCollectorMXBean;
+import java.lang.management.ManagementFactory;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+
+/**
+ * Shared helper for MDT-native vector posting lookup and approximate candidate reduction.
+ *
+ * <p>This provides the missing second-stage building block after coarse IVF cluster pruning:
+ * exact HFile lookups of cluster metadata, prefix scans over {@code P|gen|cluster|shard|},
+ * approximate RaBitQ scoring, and direct file-slice targeting from posting payload metadata.
+ */
+public final class VectorIndexMdtSearchUtils {
+
+ private static final Logger LOG = LoggerFactory.getLogger(VectorIndexMdtSearchUtils.class);
+ private static final int TOP_K_REDUCER_KEY = 0;
+ private static final Comparator<ScoredVectorPostingMatch> BEST_FIRST = Comparator
+ .comparingDouble(ScoredVectorPostingMatch::getApproxDistance)
+ .thenComparing(ScoredVectorPostingMatch::getRecordKey);
+
+ private VectorIndexMdtSearchUtils() {
+ }
+
+ public static Map<Integer, Integer> readClusterShardCounts(HoodieTableMetadata metadataTable,
+ String indexPartition,
+ int generationId,
+ Collection<Integer> clusterIds) {
+ if (clusterIds == null || clusterIds.isEmpty()) {
+ return Collections.emptyMap();
+ }
+
+ Map<Integer, Integer> shardCounts = new HashMap<>();
+ for (Integer clusterId : clusterIds) {
+ shardCounts.put(clusterId, 1);
+ }
+ return shardCounts;
+ }
+
+ public static Map<Integer, Set<String>> readClusterFileGroups(HoodieTableMetadata metadataTable,
+ String indexPartition,
+ int generationId,
+ Collection<Integer> clusterIds,
+ Collection<String> partitionPaths,
+ boolean shouldLoadInMemory) {
+ if (clusterIds == null || clusterIds.isEmpty()) {
+ return Collections.emptyMap();
+ }
+
+ return Collections.emptyMap();
+ }
+
+ public static List<VectorPostingPrefixRawKey> buildPostingPrefixes(int generationId,
+ Map<Integer, Integer> clusterShardCounts) {
+ if (clusterShardCounts == null || clusterShardCounts.isEmpty()) {
+ return Collections.emptyList();
+ }
+
+ List<Map.Entry<Integer, Integer>> clusters = new ArrayList<>(clusterShardCounts.entrySet());
+ clusters.sort(Map.Entry.comparingByKey());
+
+ List<VectorPostingPrefixRawKey> prefixes = new ArrayList<>();
+ for (Map.Entry<Integer, Integer> entry : clusters) {
+ int clusterId = entry.getKey();
+ int shardCount = Math.max(1, entry.getValue());
+ for (int shardId = 0; shardId < shardCount; shardId++) {
+ prefixes.add(new VectorPostingPrefixRawKey(generationId, clusterId, shardId));
+ }
+ }
+ return prefixes;
+ }
+
+ public static HoodieData<VectorPostingMatch> readPostingMatches(HoodieTableMetadata metadataTable,
+ String indexPartition,
+ int generationId,
+ Map<Integer, Integer> clusterShardCounts,
+ boolean shouldLoadInMemory) {
+ List<VectorPostingPrefixRawKey> postingPrefixes = buildPostingPrefixes(generationId, clusterShardCounts);
+ if (postingPrefixes.isEmpty()) {
+ return HoodieListData.eager(Collections.emptyList());
+ }
+
+ List<RawKey> rawKeys = new ArrayList<>(postingPrefixes);
+ HoodieData<VectorPostingMatch> matches = metadataTable.getRecordsByKeyPrefixes(HoodieListData.eager(rawKeys), indexPartition, shouldLoadInMemory)
+ .flatMap(record -> {
+ Option<Object> infoOpt = getVectorInfo(record);
+ if (!infoOpt.isPresent()) {
+ return Collections.<VectorPostingMatch>emptyIterator();
+ }
+
+ if (record.getData().isDeleted()) {
+ String deletedRecordKey = VectorIndexMetadataKey.postingRecordKey(record.getRecordKey());
+ return deletedRecordKey == null
+ ? Collections.<VectorPostingMatch>emptyIterator()
+ : Collections.singletonList(VectorPostingMatch.tombstone(deletedRecordKey, VectorPostingSerde.parsePostingKey(record.getRecordKey()))).iterator();
+ }
+
+ Object info = infoOpt.get();
+ int[] keyComponents = VectorPostingSerde.parsePostingKey(record.getRecordKey());
+ if (VectorPostingSerde.isPostingTombstoneInfo(info)) {
+ String deletedRecordKey = VectorIndexMetadataKey.postingRecordKey(record.getRecordKey());
+ return deletedRecordKey == null
+ ? Collections.<VectorPostingMatch>emptyIterator()
+ : Collections.singletonList(VectorPostingMatch.tombstone(deletedRecordKey, keyComponents)).iterator();
+ }
+ if (VectorPostingSerde.isPostingBlockInfo(info)) {
+ PostingBlockView view = new PostingBlockView(asPostingBlock(info));
+ List<VectorPostingMatch> blockMatches = new ArrayList<>(view.numVectors());
+ for (int vectorIndex = 0; vectorIndex < view.numVectors(); vectorIndex++) {
+ PostingBlockView.RowLocator rowLocator = view.rowLocator(vectorIndex);
+ blockMatches.add(new VectorPostingMatch(
Review Comment:
**Fixed** `c8b06883b1`. `VectorPostingMatch.fromBlockRow`, used by both call sites.
##########
hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexOptions.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.common.index.vector;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Options accepted by {@code CREATE INDEX ... USING VECTOR}.
+ *
+ * <p>The indexed column's Hudi {@code VECTOR(D[, elementType])} schema is authoritative for
+ * dimension and element type. Index options configure only the acceleration structure. DDL
+ * implementations must call {@link #validateAndNormalize(Map)} before persisting an index
+ * definition; individual parsing helpers are intentionally private so aggregate validation cannot
+ * be bypassed.
+ */
+public final class VectorIndexOptions {
+
+ public static final String METRIC = "vector.metric";
+ public static final String QUANTIZER = "vector.quantizer";
+ public static final String NUM_CLUSTERS = "vector.num_clusters";
+ public static final String MAX_ITER = "vector.max_iter";
+ public static final String RABITQ_BITS = "vector.rabitq.bits";
+ public static final String RABITQ_SEED = "vector.rabitq.seed";
+ public static final String RABITQ_ASSUME_NORMALIZED = "vector.rabitq.assume_normalized";
+ public static final String QUERY_NUM_PROBES = "vector.query.nprobes";
+ public static final String QUERY_REFINE_FACTOR = "vector.query.refine_factor";
+ public static final String QUERY_MODE = "vector.query.mode";
+ public static final String FRESHNESS_POLICY = "vector.freshness.policy";
+ public static final String STALE_LOCATOR_POLICY = "vector.stale.locator.policy";
+ public static final String FETCH_VERIFY_KEYS = "vector.fetch.verify.keys";
+
+ public static final VectorDistanceMetric DEFAULT_METRIC = VectorDistanceMetric.COSINE;
+ public static final VectorQuantizer DEFAULT_QUANTIZER = VectorQuantizer.IVF_RABITQ;
+ public static final int DEFAULT_NUM_CLUSTERS = 256;
+ public static final int DEFAULT_MAX_ITER = 20;
+ public static final int DEFAULT_RABITQ_BITS = 4;
+ public static final long DEFAULT_RABITQ_SEED = 42L;
+ public static final int DEFAULT_NUM_PROBES = 32;
+ public static final int DEFAULT_REFINE_FACTOR = 50;
+ public static final VectorQueryMode DEFAULT_QUERY_MODE = VectorQueryMode.EXACT_RERANK;
+ public static final VectorStalePolicy DEFAULT_EXACT_FRESHNESS_POLICY = VectorStalePolicy.FAIL;
+ public static final VectorStalePolicy DEFAULT_APPROXIMATE_FRESHNESS_POLICY = VectorStalePolicy.WARN;
+ public static final VectorStalePolicy DEFAULT_STALE_LOCATOR_POLICY = VectorStalePolicy.FALLBACK;
+ public static final boolean DEFAULT_FETCH_VERIFY_KEYS = true;
+
+ /** Parsed immutable view used by bootstrap and maintenance after aggregate validation. */
+ public static final class ResolvedOptions {
+ public final VectorDistanceMetric metric;
+ public final VectorQuantizer quantizer;
+ public final int numClusters;
+ public final int maxIterations;
+ public final int rabitqBits;
+ public final long rabitqSeed;
+ public final boolean assumeNormalized;
+
+ private ResolvedOptions(Map<String, String> options) {
+ this.metric = getMetric(options);
+ this.quantizer = getQuantizer(options);
+ this.numClusters = getNumClusters(options);
+ this.maxIterations = getMaxIter(options);
+ this.rabitqBits = getRaBitQBits(options);
+ this.rabitqSeed = getRaBitQSeed(options);
+ this.assumeNormalized = shouldAssumeNormalizedVectors(options);
+ }
+ }
+
+ /** Validates the complete map before exposing parsed writer-side settings. */
+ public static ResolvedOptions resolve(Map<String, String> options) {
+ return new ResolvedOptions(validateAndNormalize(options));
+ }
+
+ private static final Set<String> SUPPORTED_OPTIONS = Collections.unmodifiableSet(
+ new HashSet<>(Arrays.asList(
+ METRIC,
+ QUANTIZER,
+ NUM_CLUSTERS,
+ MAX_ITER,
+ RABITQ_BITS,
+ RABITQ_SEED,
+ RABITQ_ASSUME_NORMALIZED,
+ QUERY_NUM_PROBES,
+ QUERY_REFINE_FACTOR,
+ QUERY_MODE,
+ FRESHNESS_POLICY,
+ STALE_LOCATOR_POLICY,
+ FETCH_VERIFY_KEYS)));
+
+ private VectorIndexOptions() {
+ }
+
+ /**
+ * Validates the complete option map and returns canonical values for persistence.
+ *
+ * <p>The returned map contains every supported option, including explicit defaults. Unknown,
+ * retired, misspelled, and invalid options are rejected instead of being silently ignored.
+ */
+ public static Map<String, String> validateAndNormalize(Map<String, String> options) {
+ Set<String> unknownOptions = new HashSet<>(options.keySet());
+ unknownOptions.removeAll(SUPPORTED_OPTIONS);
+ if (!unknownOptions.isEmpty()) {
+ throw new IllegalArgumentException("Unsupported vector index options: " + unknownOptions);
+ }
+
+ VectorDistanceMetric metric = getMetric(options);
+ VectorQuantizer quantizer = getQuantizer(options);
+ int numClusters = getNumClusters(options);
+ int maxIter = getMaxIter(options);
+ int bits = getRaBitQBits(options);
+ long seed = getRaBitQSeed(options);
+ boolean assumeNormalized = shouldAssumeNormalizedVectors(options);
+ int numProbes = getNumProbes(options);
+ int refineFactor = getRefineFactor(options);
+ VectorQueryMode queryMode = getQueryMode(options);
+ VectorStalePolicy freshnessPolicy = getFreshnessPolicy(options);
+ VectorStalePolicy staleLocatorPolicy = getStaleLocatorPolicy(options);
+ boolean verifyFetchKeys = shouldVerifyFetchKeys(options);
+
+ if (numProbes > numClusters) {
+ throw new IllegalArgumentException(
+ "Option '" + QUERY_NUM_PROBES + "' must not exceed '" + NUM_CLUSTERS + "': "
+ + numProbes + " > " + numClusters);
+ }
+
+ Map<String, String> normalized = new LinkedHashMap<>();
+ normalized.put(METRIC, metric.name().toLowerCase(Locale.ROOT));
+ normalized.put(QUANTIZER, quantizer.name());
+ normalized.put(NUM_CLUSTERS, String.valueOf(numClusters));
+ normalized.put(MAX_ITER, String.valueOf(maxIter));
+ normalized.put(RABITQ_BITS, String.valueOf(bits));
+ normalized.put(RABITQ_SEED, String.valueOf(seed));
+ normalized.put(RABITQ_ASSUME_NORMALIZED, String.valueOf(assumeNormalized));
+ normalized.put(QUERY_NUM_PROBES, String.valueOf(numProbes));
+ normalized.put(QUERY_REFINE_FACTOR, String.valueOf(refineFactor));
+ normalized.put(QUERY_MODE, queryMode.name().toLowerCase(Locale.ROOT));
+ normalized.put(FRESHNESS_POLICY, freshnessPolicy.name().toLowerCase(Locale.ROOT));
+ normalized.put(STALE_LOCATOR_POLICY, staleLocatorPolicy.name().toLowerCase(Locale.ROOT));
+ normalized.put(FETCH_VERIFY_KEYS, String.valueOf(verifyFetchKeys));
+ return Collections.unmodifiableMap(normalized);
+ }
+
+ public static VectorDistanceMetric getMetric(Map<String, String> options) {
+ String value = getOption(options, METRIC, DEFAULT_METRIC.name());
+ try {
+ return VectorDistanceMetric.fromString(value);
+ } catch (IllegalArgumentException e) {
+ throw unsupportedValue(METRIC, value, e);
+ }
+ }
+
+ private static VectorQuantizer getQuantizer(Map<String, String> options) {
+ String value = getOption(options, QUANTIZER, DEFAULT_QUANTIZER.name());
+ try {
+ return VectorQuantizer.fromString(value);
+ } catch (IllegalArgumentException e) {
+ throw unsupportedValue(QUANTIZER, value, e);
+ }
+ }
+
+ private static int getNumClusters(Map<String, String> options) {
+ return getPositiveInt(options, NUM_CLUSTERS, DEFAULT_NUM_CLUSTERS);
+ }
+
+ private static int getMaxIter(Map<String, String> options) {
+ return getPositiveInt(options, MAX_ITER, DEFAULT_MAX_ITER);
+ }
+
+ private static int getRaBitQBits(Map<String, String> options) {
+ int bits = getInt(options, RABITQ_BITS, DEFAULT_RABITQ_BITS);
+ if (bits < 1 || bits > 8) {
+ throw new IllegalArgumentException(
+ "Option '" + RABITQ_BITS + "' must be between 1 and 8: " + bits);
+ }
+ return bits;
+ }
+
+ private static long getRaBitQSeed(Map<String, String> options) {
+ String value = getOption(options, RABITQ_SEED, String.valueOf(DEFAULT_RABITQ_SEED));
+ try {
+ return Long.parseLong(value);
+ } catch (NumberFormatException e) {
+ throw invalidNumber(RABITQ_SEED, value, e);
+ }
+ }
+
+ private static boolean shouldAssumeNormalizedVectors(Map<String, String> options) {
+ return getBoolean(options, RABITQ_ASSUME_NORMALIZED, false);
+ }
+
+ public static int getNumProbes(Map<String, String> options) {
+ return getPositiveInt(options, QUERY_NUM_PROBES, DEFAULT_NUM_PROBES);
+ }
+
+ public static int getRefineFactor(Map<String, String> options) {
+ return getPositiveInt(options, QUERY_REFINE_FACTOR, DEFAULT_REFINE_FACTOR);
+ }
+
+ private static VectorQueryMode getQueryMode(Map<String, String> options) {
+ String value = getOption(options, QUERY_MODE, DEFAULT_QUERY_MODE.name());
+ try {
+ return VectorQueryMode.fromString(value);
+ } catch (IllegalArgumentException e) {
+ throw unsupportedValue(QUERY_MODE, value, e);
+ }
+ }
+
+ public static VectorStalePolicy getFreshnessPolicy(Map<String, String> options) {
+ VectorStalePolicy defaultPolicy = getQueryMode(options) == VectorQueryMode.APPROXIMATE
+ ? DEFAULT_APPROXIMATE_FRESHNESS_POLICY : DEFAULT_EXACT_FRESHNESS_POLICY;
+ return getStalePolicy(options, FRESHNESS_POLICY, defaultPolicy);
+ }
+
+ public static VectorStalePolicy getStaleLocatorPolicy(Map<String, String> options) {
+ return getStalePolicy(options, STALE_LOCATOR_POLICY, DEFAULT_STALE_LOCATOR_POLICY);
+ }
+
+ private static VectorStalePolicy getStalePolicy(
+ Map<String, String> options, String key, VectorStalePolicy defaultPolicy) {
+ String value = getOption(options, key, defaultPolicy.name());
+ try {
+ return VectorStalePolicy.fromString(value);
+ } catch (IllegalArgumentException e) {
+ throw unsupportedValue(key, value, e);
+ }
+ }
+
+ public static boolean isApproximateSearchMode(Map<String, String> options) {
+ return getQueryMode(options) == VectorQueryMode.APPROXIMATE;
+ }
+
+ /** The current posting format uses symmetric scoring. */
+ public static boolean isRaBitQAsymmetricScoring(Map<String, String> options) {
+ return false;
Review Comment:
**Fixed** `3cde5283ea`. Replaced with constants / removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-6448261785
Compatibility question before we retain any vector-specific freshness machinery: @cvinoth, is the release/table format carrying the vector index intended to require a table-version bump that fences older writers? Current master routes configured/inflight metadata partition paths through MetadataPartitionType.fromPartitionPath, which rejects unknown vector_index_* paths, so an older writer that observes the partition fails closed. We have not yet established the behavior of a pre-existing writer with cached pre-index table configuration. We want to reuse Hudi RLI/SI compatibility semantics rather than add a per-commit frontier for a writer that is already fenced. Could you confirm the intended writer compatibility/table-version contract?
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "hudi-bot (via GitHub)" <gi...@apache.org>.
hudi-bot commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-5860422709
<!--
Meta data
{
"version" : 1,
"metaDataEntries" : [ {
"hash" : "4204280759551a140914b8a393d0420206489542",
"status" : "PENDING",
"url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=17318",
"triggerID" : "4204280759551a140914b8a393d0420206489542",
"triggerType" : "PUSH"
} ]
}-->
## CI report:
* 4204280759551a140914b8a393d0420206489542 Azure: [PENDING](https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=17318)
<details>
<summary>Bot commands</summary>
@hudi-bot supports the following commands:
- `@hudi-bot run azure` re-run the last Azure build
</details>
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "hudi-bot (via GitHub)" <gi...@apache.org>.
hudi-bot commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-5860379912
<!--
Meta data
{
"version" : 1,
"metaDataEntries" : [ {
"hash" : "4204280759551a140914b8a393d0420206489542",
"status" : "UNKNOWN",
"url" : "TBD",
"triggerID" : "4204280759551a140914b8a393d0420206489542",
"triggerType" : "PUSH"
} ]
}-->
## CI report:
* 4204280759551a140914b8a393d0420206489542 UNKNOWN
<details>
<summary>Bot commands</summary>
@hudi-bot supports the following commands:
- `@hudi-bot run azure` re-run the last Azure build
</details>
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "hudi-bot (via GitHub)" <gi...@apache.org>.
hudi-bot commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-5861183452
<!--
Meta data
{
"version" : 1,
"metaDataEntries" : [ {
"hash" : "4204280759551a140914b8a393d0420206489542",
"status" : "SUCCESS",
"url" : "https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=17318",
"triggerID" : "4204280759551a140914b8a393d0420206489542",
"triggerType" : "PUSH"
} ]
}-->
## CI report:
* 4204280759551a140914b8a393d0420206489542 Azure: [SUCCESS](https://dev.azure.com/apachehudi/a1a51da7-8592-47d4-88dc-fd67bed336bb/_build/results?buildId=17318)
<details>
<summary>Bot commands</summary>
@hudi-bot supports the following commands:
- `@hudi-bot run azure` re-run the last Azure build
</details>
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-6457703456
@yihua @vinishjail97 The review replies are updated, and the synchronous-maintenance follow-up is published at https://github.com/chrevanthreddy/hudi/pull/5 (head `0674d3d794`). Ready for another review pass.
The follow-up removes per-write manifest-frontier advancement: postings and source-instant audit markers are maintained together, while strict read validation is opt-in/default off. Policy precedence and required bootstrap-baseline validation are covered by focused tests. Base remains bootstrap + read path; maintenance is in the stacked PR.
Latest vector-only Spark 4.0 validation: common 13/13 and planner 9/9 passed; combined E2E finished 6 tests with one retry-dependent inflight-bootstrap pass. That concurrent config-read exception is explicitly recorded in the stacked PR; the separate shared config-reader candidate fix is excluded. Existing per-thread replies contain the implementation details. Thanks for taking another look.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "vinothchandar (via GitHub)" <gi...@apache.org>.
vinothchandar commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-5921496452
is this the latest PR.. like the new one for us to review @chrevanthreddy ..
@vinishjail97 is going to work with you closely and get this over.. test this within Onehouse as well parallely..
@yihua and I can help reviews as needed
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4384297540
##########
hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkVectorIndexBootstrap.java:
##########
@@ -0,0 +1,531 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.metadata;
+
+import org.apache.hudi.common.data.HoodieData;
+import org.apache.hudi.common.index.vector.PostingBlockBuilder;
+import org.apache.hudi.common.index.vector.RaBitQEncoder;
+import org.apache.hudi.common.index.vector.VectorDistanceMetric;
+import org.apache.hudi.common.index.vector.VectorIndexBootstrapUtils;
+import org.apache.hudi.common.index.vector.VectorIndexOptions;
+import org.apache.hudi.common.model.HoodieIndexDefinition;
+import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.schema.HoodieSchema;
+import org.apache.hudi.data.HoodieJavaRDD;
+import org.apache.hudi.spark.index.vector.TwoLevelKMeansBootstrap$;
+
+import lombok.extern.slf4j.Slf4j;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.broadcast.Broadcast;
+import org.apache.spark.sql.SparkSession;
+
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import scala.Tuple2;
+
+/**
+ * Two-phase vector index bootstrap for Spark.
+ *
+ * <p>Phase 1: Sample vectors → train centroids with Spark ML KMeans. Training
+ * still uses a bounded sample (FAISS, ScaNN, and Milvus all use sample-based
+ * centroid training for IVF), but the iterative KMeans work stays distributed
+ * instead of collecting the sample to the driver.
+ *
+ * <p>Phase 2: Broadcast centroids → single-pass {@code mapPartitions} that
+ * assigns each vector to its nearest centroid, computes RaBitQ binary codes,
+ * and emits posting metadata records with per-candidate locators. Cluster
+ * statistics (populations, shard counts) are gathered via a lightweight
+ * {@code countByKey} before the main encoding pass.
+ *
+ * <p>Compared to the previous driver-local KMeans approach:
+ * <ul>
+ * <li>No reflection — direct {@link HoodieMetadataPayload} method calls</li>
+ * <li>Proper {@code unpersist()} in try/finally</li>
+ * <li>generation ordinal allocated from existing manifest keys</li>
+ * </ul>
+ */
+@Slf4j
+public final class SparkVectorIndexBootstrap {
+
+ /** Hard cap for percentile-driven training samples. */
+ private static final int MAX_PERCENTILE_TRAINING_SAMPLE = 10_000_000;
+
+ /** Absolute minimum sample floor for production index quality. */
+ private static final int MIN_TRAINING_SAMPLE = 1_000_000;
+
+ /** Minimum sample multiplier per cluster. FAISS recommends 256×K minimum. */
+ private static final int MIN_SAMPLE_PER_CLUSTER = 256;
+
+ /** Lower bound for percentile-driven sampling: 0.5% of total vectors. */
+ private static final double MIN_SAMPLE_PERCENT = 0.005d;
+
+ /** Upper bound for percentile-driven sampling: 1% of total vectors. */
+ private static final double MAX_SAMPLE_PERCENT = 0.01d;
+
+ private SparkVectorIndexBootstrap() {
+ }
+
+ /**
+ * Run the two-phase vector index bootstrap and return all MDT records.
+ *
+ * @param jsc Spark context
+ * @param vectorRows RDD of (recordKey, partitionPath, fileId, vectorBytes) — already
+ * read from base-table file slices
+ * @param indexDef the index definition
+ * @param vectorType resolved vector element type from the table schema
+ * @param generation vector generation ordinal
+ * @param lastUpdatedTs timestamp for MDT bookkeeping
+ * @return all metadata records to write into the vector index MDT partition
+ */
+ public static HoodieData<HoodieRecord> bootstrap(JavaSparkContext jsc,
+ JavaRDD<VectorRow> vectorRows,
+ HoodieIndexDefinition indexDef,
+ HoodieSchema.Vector.VectorElementType vectorType,
+ int dimension,
+ int generation,
+ long lastUpdatedTs) {
+ Map<String, String> options = indexDef.getIndexOptions();
+ String indexName = indexDef.getIndexName();
+ String vectorColumn = indexDef.getSourceFields().isEmpty() ? "" : indexDef.getSourceFields().get(0);
+ VectorIndexOptions.ResolvedOptions resolvedOptions = VectorIndexOptions.resolve(options);
+ int numClusters = resolvedOptions.numClusters;
+ VectorDistanceMetric metric = resolvedOptions.metric;
+
+ String quantizerType = resolvedOptions.quantizer.name();
+ long quantizerSeed = resolvedOptions.rabitqSeed;
+ int rabitqBits = resolvedOptions.rabitqBits;
+ boolean assumeNormalized = resolvedOptions.assumeNormalized;
+ boolean storeInMdt = true;
+ int targetRowsPerShard = 100_000;
Review Comment:
**Fixed** `3cde5283ea`. Flag and dead branches removed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-6431609261
Thanks @vinishjail97, all of your stacked fixes are folded in (vector_index type, bootstrap instant in the manifest,
coarse-cell clamp, MSB-first extended planes, metric/bits DDL validation with the l2 default, bits=1 at non-64
dimensions, and `TestHoodieVectorIndexSearch`). Of your follow-ups: incremental updates → https://github.com/chrevanthreddy/hudi/pull/5;
partitioned RLI → fixed here (thread above); chunked centroids and bits=1 cosine/dot stay open as follow-ups.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "yihua (via GitHub)" <gi...@apache.org>.
yihua commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4214245877
##########
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:
@chrevanthreddy do you plan to have a follow-up PR for the implementation of the synchronous update of vector 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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-6431609632
Thanks @yihua, @vinishjail97 and the bot for the reviews. The changes and follow-up scope are summarized below; per-thread replies are included below.
**What changed since the last review**
- Partitioned RLI finalists are batched per RLI file group (yihua).
- Bootstrap frontier follows the async-indexer cutoff (in-flight writers excluded).
- Exact search supports all column types, reads MOR log rows on executors, and returns current MOR column values.
- Query option validation, catalog-qualified names, NULL TVF args, bounded caches, plus dead code / duplication cleanup.
**Scope / stack**
1. This PR: schema, RaBitQ, bootstrap, read path.
2. https://github.com/chrevanthreddy/hudi/pull/4: metadata-key routing for vector lookups (1B: first query 734s → 28.5s; repeat-query p50 17.3s → 7.0s approx, 24.5s → 14.8s exact; recall unchanged).
3. https://github.com/chrevanthreddy/hudi/pull/5: synchronous index maintenance (`buildUpdate`), with the lifecycle validation; depends on #19801.
The shared-code changes include vector schema/payload support, index scheduling/configuration integration, and reuse of the existing Parquet row-read configuration. The payload slow path keeps its positional contract and adds a schema guard. MOR row positions come from the base file using the existing key/position utility, avoiding the proposed shared row-index ordinal fix. The separate routing PR changes `HoodieBackedTableMetadata`, restricted to vector partitions.
**Validation**
- 1B COW (BigANN, 16384 clusters, 4-bit RaBitQ, nprobes=32, refine=50): approx recall 0.865 / exact 0.95, p50 6.8s / 12.5s with https://github.com/chrevanthreddy/hudi/pull/4 (20 queries × 2 passes).
- Lifecycle with exact == brute force before and after upsert/delete/clustering at 10M COW+MOR and 1B COW.
- Local Spark 4.0 validation: all three stacked branches passed builds/style and the selected unit/end-to-end suites. Base/routing: 162 common, 12 hadoop-common, 40 client-common, 7 spark-client, 110 Spark tests each. Update: 164 / 12 / 50 / 26 / 114, all zero failures/errors. This run predates only test-output/comment cleanup and an EOF whitespace change; the two cleaned-up residual tests passed separately.
- Routing's focused regression suite passed 8/8 separately (the original broad selector omitted it).
- Spark 3.5 builds/style passed on all branches. Common tests: 228 / 236 / 238; hadoop-common: 12 each. Base/routing E2E: 110 each; update E2E: 114 with one retry-dependent in-flight-writer flake (passed three subsequent independent runs with retries disabled; root cause not established). Unmodified client selections encountered missing Hadoop shaded runtime classes. Matching-runtime diagnostic reruns passed client-common 45 / 45 / 55 and spark-client 13 / 13 / 35 on base/routing/update respectively; these are classpath-adjusted results, not unmodified-build passes.
- Supplemental update Spark 4.0 tests passed: common 79, client 17, streaming 9, Spark-common 16. Flink 2.2 build passed; handler 2 passed and coordinator/event buffers 72 passed with 1 skipped. All results are local, not remote CI. Validation caveats are retained here and in the stacked PR descriptions.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4384294616
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/PositionalParquetFetcher.scala:
##########
@@ -0,0 +1,587 @@
+/*
+ * 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.spark.sql.hudi.analysis
+
+import org.apache.hudi.common.index.vector.{ScoredVectorPostingMatch, VectorIndexMdtSearchUtils, VectorPostingMatch}
+import org.apache.hudi.common.model.HoodieRecord
+import org.apache.hudi.common.schema.HoodieSchema
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileSystem, Path}
+import org.apache.parquet.column.page.PageReadStore
+import org.apache.parquet.hadoop.ParquetFileReader
+import org.apache.parquet.hadoop.metadata.{BlockMetaData, ColumnChunkMetaData, ParquetMetadata}
+import org.apache.parquet.internal.column.columnindex.OffsetIndex
+import org.apache.parquet.internal.filter2.columnindex.RowRanges
+import org.apache.parquet.io.{ColumnIOFactory, RecordReader}
+import org.apache.parquet.io.api.{Binary, Converter, GroupConverter, PrimitiveConverter, RecordMaterializer}
+import org.apache.parquet.schema.MessageType
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction.DistanceMetric
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+import org.slf4j.LoggerFactory
+
+import java.util.PrimitiveIterator
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+/**
+ * Fetches exact vector-search rerank candidates from Parquet by physical row position.
+ *
+ * The fetcher uses Parquet page indexes to turn candidate row positions into row ranges, so
+ * the reader can skip non-candidate pages instead of scanning whole Hudi file slices.
+ */
+private[analysis] final class PositionalParquetFetcher(
+ conf: Configuration,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Double],
+ metric: DistanceMetric.Value,
+ outputSchema: StructType,
+ failOnMissingOffsetIndex: Boolean,
+ vectorOnlyProjection: Boolean,
+ verifyKeys: Boolean) extends Serializable {
+
+ import PositionalParquetFetcher._
+
+ private val distanceOrdinal = outputSchema.fieldIndex(HoodieVectorSearchPlanBuilder.DISTANCE_COL)
+ private val recordKeyOrdinal = fieldOrdinal(HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ private val partitionPathOrdinal = fieldOrdinal(HoodieRecord.PARTITION_PATH_METADATA_FIELD)
+ private val fileGroupIdOrdinal = fieldOrdinal(HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL)
+ private val scorer = new VectorExactScorer(vectorSchema, queryVector, metric)
+ private var scoreMs: Long = 0L
+
+ def fetch(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch]): FetchResult = {
+ val startNs = System.nanoTime()
+ if (candidates.isEmpty) {
+ FetchResult(Iterator.empty, Seq.empty, FetchMetrics(elapsedMs = elapsedMs(startNs)))
+ } else {
+ fetchNonEmpty(filePath, candidates, startNs)
+ }
+ }
+
+ private def fetchNonEmpty(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch],
+ startNs: Long): FetchResult = {
+ val path = new Path(filePath)
+ configureRandomRead(conf)
+ val fs = path.getFileSystem(conf)
+ val fileLen = fs.getFileStatus(path).getLen
+ val fileCacheKey = s"$filePath#$fileLen"
+ val stats = Option(FileSystem.getStatistics(path.toUri.getScheme, fs.getClass))
+ val physicalBefore = stats.map(_.getBytesRead).getOrElse(0L)
+
+ val footerStartNs = System.nanoTime()
+ val footerResult = footer(conf, path, fileCacheKey)
+ val footerMs = elapsedMs(footerStartNs)
+ val metadata = footerResult.metadata
+ val fileSchema = metadata.getFileMetaData.getSchema
+ val requestedParquetSchema = requestedSchema(fileSchema)
+
+ val rowGroups = metadata.getBlocks.asScala.toIndexedSeq
+ val rowGroupStarts = cumulativeRowGroupStarts(rowGroups)
+ val byRowGroup: Map[Int, Seq[ScoredVectorPostingMatch]] =
+ candidates.groupBy(candidate => rowGroupOrdinal(rowGroupStarts, rowGroups, rowPosition(candidate)))
+
+ val rows = mutable.ArrayBuffer.empty[InternalRow]
+ val keyMismatches = mutable.ArrayBuffer.empty[ScoredVectorPostingMatch]
+ var rowGroupsSelected = 0
+ var pagesInSelectedRowGroups = 0L
+ var pagesSelected = 0L
+ var offsetIndexHits = 0L
+ var offsetIndexMissing = 0L
+ var pageBytesFetched = 0L
+ var rangedGets = 0L
+ var rowsDecoded = 0L
+ var rowsMaterialized = 0L
+ var offsetIndexMs = 0L
+ var decodeMs = 0L
+ scoreMs = 0L
+
+ val reader = new ParquetFileReader(conf, path, metadata)
+ try {
+ reader.setRequestedSchema(requestedParquetSchema)
+ byRowGroup.toSeq.sortBy(_._1).foreach { case (rgOrdinal, rgCandidates) =>
+ val block = rowGroups(rgOrdinal)
+ val rowGroupStart = rowGroupStarts(rgOrdinal)
+ val embeddingChunk = findColumn(block, embeddingCol)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet file '$filePath' is missing vector column '$embeddingCol'."))
+
+ val offsetStartNs = System.nanoTime()
+ val embeddingOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, embeddingChunk)
+ offsetIndexMs += elapsedMs(offsetStartNs)
+ val relativeRows = rgCandidates.map(candidate => rowPosition(candidate) - rowGroupStart).sorted
+ val candidateRows = relativeRows.toSet
+ val candidateByRelativeRow = rgCandidates
+ .map(candidate => rowPosition(candidate) - rowGroupStart -> candidate)
+ .toMap
+
+ if (embeddingOffsetIndex == null && failOnMissingOffsetIndex) {
+ offsetIndexMissing += 1
+ throw new IllegalStateException(
+ s"Parquet file '$filePath' has no OffsetIndex for vector column '$embeddingCol'; positional exact fetch cannot honor pages.")
+ }
+
+ rowGroupsSelected += 1
+ val (pageStore, rowRanges) = if (embeddingOffsetIndex == null) {
+ offsetIndexMissing += 1
+ LOG.warn(
+ "Parquet file {} rowGroup={} has no OffsetIndex for vector column {}; falling back to row-group fetch.",
+ filePath,
+ Int.box(rgOrdinal),
+ embeddingCol)
+ pagesInSelectedRowGroups += 0L
+ pagesSelected += 0L
+ rangedGets += 1L
+ pageBytesFetched += requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .map(_.getTotalSize)
+ .sum
+ (reader.readRowGroup(rgOrdinal), RowRanges.createSingle(block.getRowCount))
+ } else {
+ offsetIndexHits += 1
+ pagesInSelectedRowGroups += embeddingOffsetIndex.getPageCount
+ val selectedVectorPages = selectedPages(embeddingOffsetIndex, block.getRowCount, relativeRows)
+ pagesSelected += selectedVectorPages.size
+ val selectedRanges = requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .flatMap { chunk =>
+ val oiStartNs = System.nanoTime()
+ val chunkOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, chunk)
+ offsetIndexMs += elapsedMs(oiStartNs)
+ if (chunkOffsetIndex == null) {
+ offsetIndexMissing += 1
+ Seq.empty
+ } else {
+ offsetIndexHits += 1
+ selectedPages(chunkOffsetIndex, block.getRowCount, relativeRows).map { pageOrdinal =>
+ val start = chunkOffsetIndex.getOffset(pageOrdinal)
+ PageRange(start, start + chunkOffsetIndex.getCompressedPageSize(pageOrdinal))
+ }
+ }
+ }
+ val mergedRanges = mergeRanges(selectedRanges.sortBy(_.start).toSeq)
+ rangedGets += mergedRanges.size
+ pageBytesFetched += mergedRanges.iterator.map(range => range.end - range.start).sum
+ val ranges = RowRanges.create(block.getRowCount, primitiveIntIterator(selectedVectorPages), embeddingOffsetIndex)
+ (reader.readFilteredRowGroup(rgOrdinal, ranges), ranges)
+ }
+
+ val decodeStartNs = System.nanoTime()
+ val recordReader = parquetRecordReader(fileSchema, requestedParquetSchema, pageStore)
+ val rowIterator = rowRanges.iterator()
+ while (rowIterator.hasNext) {
+ val relativeRow = rowIterator.nextLong()
+ val row = recordReader.read()
+ rowsDecoded += 1
+ candidateByRelativeRow.get(relativeRow).foreach { candidate =>
+ if (verifyKeys && decodedRecordKey(row) != candidate.getRecordKey) {
+ keyMismatches += candidate
+ } else {
+ fillSyntheticFields(row, candidate)
+ rows += row
+ rowsMaterialized += 1
+ }
+ }
+ }
+ decodeMs += elapsedMs(decodeStartNs)
+ }
+ } finally {
+ reader.close()
+ }
+
+ val physicalAfter = stats.map(_.getBytesRead).getOrElse(physicalBefore)
+ val metrics = FetchMetrics(
+ rerankCandidates = candidates.size,
+ rowGroupsTotal = rowGroups.size,
+ rowGroupsSelected = rowGroupsSelected,
+ pagesInSelectedRowGroups = pagesInSelectedRowGroups,
+ pagesSelected = pagesSelected,
+ offsetIndexHits = offsetIndexHits,
+ offsetIndexMissing = offsetIndexMissing,
+ rangedGets = rangedGets,
+ pageBytesFetched = pageBytesFetched,
+ rowsDecoded = rowsDecoded,
+ rowsMaterialized = rowsMaterialized,
+ fetchKeyMismatch = keyMismatches.size,
+ physicalBytesRead = math.max(0L, physicalAfter - physicalBefore),
+ footerReads = if (footerResult.cacheHit) 0 else 1,
+ footerCacheHits = if (footerResult.cacheHit) 1 else 0,
+ footerMs = footerMs,
+ offsetIndexMs = offsetIndexMs,
+ fetchWaitMs = 0L,
+ decodeMs = decodeMs,
+ scoreMs = scoreMs,
+ elapsedMs = elapsedMs(startNs))
+ FetchResult(rows.iterator, keyMismatches.toSeq, metrics)
+ }
+
+ private def parquetRecordReader(
+ fileSchema: MessageType,
+ requestedParquetSchema: MessageType,
+ pageStore: PageReadStore): RecordReader[InternalRow] = {
+ val materializer = new InternalRowMaterializer(requestedParquetSchema, outputSchema, embeddingCol, distanceOrdinal)
+ new ColumnIOFactory().getColumnIO(requestedParquetSchema, fileSchema)
+ .getRecordReader(pageStore, materializer)
+ }
+
+ private def requestedSchema(fileSchema: MessageType): MessageType = {
+ val parquetFields =
+ if (vectorOnlyProjection) {
+ mutable.ArrayBuffer.empty[org.apache.parquet.schema.Type]
+ } else {
+ outputSchema.fields
+ .filterNot(field => isSyntheticOutputField(field.name))
+ .map(field => fieldByName(fileSchema, field.name))
+ .toBuffer[org.apache.parquet.schema.Type]
+ }
+ if ((verifyKeys || vectorOnlyProjection) && !parquetFields.exists(_.getName == HoodieRecord.RECORD_KEY_METADATA_FIELD)) {
+ parquetFields += fieldByName(fileSchema, HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ }
+ if (!parquetFields.exists(_.getName == embeddingCol)) {
+ parquetFields += fieldByName(fileSchema, embeddingCol)
+ }
+ new MessageType(fileSchema.getName, parquetFields.asJava)
+ }
+
+ private def isSyntheticOutputField(name: String): Boolean =
+ name == HoodieVectorSearchPlanBuilder.DISTANCE_COL ||
+ name == HoodieRecord.RECORD_KEY_METADATA_FIELD ||
+ name == HoodieRecord.PARTITION_PATH_METADATA_FIELD ||
+ name == HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL
+
+ private def fieldByName(schema: MessageType, name: String): org.apache.parquet.schema.Type =
+ schema.getFields.asScala.find(_.getName == name)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet schema is missing requested field '$name'."))
+
+ private def rowPosition(candidate: ScoredVectorPostingMatch): Long = {
+ val location = candidate.getLocation
+ if (location == null) {
+ throw new IllegalArgumentException(
+ s"Vector exact positional fetch requires a row-position locator for record ${candidate.getRecordKey}.")
+ }
+ location.getPosition
+ }
+
+ private def decodedRecordKey(row: InternalRow): String =
+ recordKeyOrdinal.flatMap(ordinal => Option(row.getUTF8String(ordinal))).map(_.toString).orNull
+
+ private def fillSyntheticFields(row: InternalRow, candidate: ScoredVectorPostingMatch): Unit = {
+ val liveLocation = Option(candidate.getLocation)
+ recordKeyOrdinal.foreach(row.update(_, UTF8String.fromString(candidate.getRecordKey)))
+ partitionPathOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getPartitionPath).orElse(Option(candidate.getPartitionPath)).getOrElse(""))))
+ fileGroupIdOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getFileId).orElse(Option(candidate.getFileGroupId)).getOrElse(""))))
+ }
+
+ private def fieldOrdinal(name: String): Option[Int] =
+ outputSchema.fields.indexWhere(_.name == name) match {
+ case -1 => None
+ case ordinal => Some(ordinal)
+ }
+
+ private def scoreVector(bytes: Array[Byte]): Double = scorer.scoreBytes(bytes)
+
+ private final class InternalRowMaterializer(
+ requestedParquetSchema: MessageType,
+ outputSchema: StructType,
+ embeddingCol: String,
+ distanceOrdinal: Int) extends RecordMaterializer[InternalRow] {
+
+ private val values = new Array[Any](outputSchema.length)
+ private var current: InternalRow = _
+ private var vectorBytes: Array[Byte] = _
+
+ private val root = new GroupConverter {
+ private val converters = requestedParquetSchema.getFields.asScala.map(converterFor).toIndexedSeq
+
+ override def getConverter(fieldIndex: Int): Converter =
+ converters(fieldIndex)
+
+ override def start(): Unit = {
+ var idx = 0
+ while (idx < values.length) {
+ values(idx) = null
+ idx += 1
+ }
+ vectorBytes = null
+ }
+
+ override def end(): Unit = {
+ if (vectorBytes == null) {
+ throw new IllegalStateException(s"Vector column '$embeddingCol' was not decoded for positional exact fetch.")
Review Comment:
**Fixed**. A live row with a NULL embedding is skipped (counted as `nullEmbeddingsSkipped`), matching brute force.
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorExactSearchPlanner.scala:
##########
@@ -0,0 +1,390 @@
+/*
+ * 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.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.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, 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}
+
+/** Exact reranking and file-slice resolution for the IVF/RaBitQ search plan. */
+private[analysis] object VectorExactSearchPlanner extends SparkAdapterSupport {
+ import HoodieVectorSearchPlanBuilder._
+ import IvfRaBitQMdtSearchAlgorithm.{isLogResidentCandidate, isPartitionedRecordIndex, toVectorMetric, VectorSearchPlanContext}
+
+ private val LOG = LoggerFactory.getLogger(getClass)
+
+ def buildRaBitQExactCandidateDf(
+ spark: SparkSession,
+ ctx: VectorSearchPlanContext,
+ outputSchema: StructType,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Float],
+ exactQueryVector: Array[Double],
+ k: Int): DataFrame = {
+ val planStartNs = System.nanoTime()
+ val storageConf = HadoopFSUtils.getStorageConf(spark.sessionState.newHadoopConf())
+ val cacheLoadStartNs = 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 topClusters = cache.findTopClusters(queryVector, numProbes, metric)
+ val topClusterSeq = topClusters.toSeq
+ val cacheLoadMs = elapsedMs(cacheLoadStartNs)
+
+ // Shard counts from cache — no extra MDT IO.
+ val shardCounts = cache.getShardCounts(topClusters)
+ val shardCountsScala = shardCounts.asScala.map { case (clusterId, count) =>
+ clusterId.intValue() -> count.intValue()
+ }.toMap
+ val postingPrefixCount = VectorIndexMdtSearchUtils.buildPostingPrefixes(cache.getGenerationId, shardCounts).size()
+ LOG.info(
+ s"[vector_search][plan] basePath=${ctx.basePath} indexPartition=${ctx.indexPartition} cacheLoadMs=$cacheLoadMs " +
+ s"generationId=${cache.getGenerationId} dimension=${cache.getDimension} numClusters=${cache.numClusters()} " +
+ s"numProbes=$numProbes refineFactor=${VectorIndexOptions.getRefineFactor(mergedOptions.asJava)} residualEncoding=$residualEncoding " +
+ s"topClusters=${summarizeInts(topClusterSeq)} shardCounts=${summarizeShardCounts(shardCountsScala)} " +
+ s"postingPrefixes=$postingPrefixCount")
+ if (shardCounts.isEmpty) {
+ sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+ } else {
+ val refineK = math.max(k, k * VectorIndexOptions.getRefineFactor(mergedOptions.asJava))
+ val emptyDf = sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+
+ 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, refineK)
+ val topCandidateMaterializeStartNs = System.nanoTime()
+ val materializedCandidates = try {
+ topCandidates.collectAsList().asScala.toSeq
+ } finally {
+ topCandidates.unpersistWithDependencies()
+ }
+ val topCandidateList = materializedCandidates
+ .filter(candidate => candidate.getFileGroupId != null && candidate.getFileGroupId.nonEmpty)
+ val candidatesWithoutLocator = materializedCandidates.size - topCandidateList.size
+ LOG.info(
+ s"[vector_search][stage][materialize_top_candidates] refineK=$refineK rerankCandidates=${materializedCandidates.size} " +
+ s"candidatesWithLocators=${topCandidateList.size} candidatesWithoutLocators=$candidatesWithoutLocator " +
+ s"distinctFileGroups=${topCandidateList.map(_.getFileGroupId).distinct.size} elapsedMs=${elapsedMs(topCandidateMaterializeStartNs)}")
+ if (candidatesWithoutLocator > 0) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search found $candidatesWithoutLocator rerank candidates without file-group locators. " +
+ s"Exact rerank requires per-candidate posting locators.")
+ }
+
+ // RFC-109 RLI finalist arbiter. Exact mode uses refreshed RLI locations for every live
+ // candidate: SERVE retains positional trust; STALE is resolved by key in the live base file.
+ val arbiterEnabled = VectorIndexOptions.isFinalistArbiterEnabled(mergedOptions.asJava)
+ val serveCandidates: Seq[ScoredVectorPostingMatch] =
+ if (arbiterEnabled) {
+ val arb = VectorIndexMdtSearchUtils.arbitrateMaterializedFinalists(
+ ctx.metadataTable, topCandidateList.asJava, isPartitionedRecordIndex(ctx))
+ val staleFailPolicy = VectorIndexOptions.isStaleLocatorPolicyFail(mergedOptions.asJava)
+ LOG.info(
+ s"[vector_search][exact][arbiter] arbiterExclusions{stale=${arb.staleCount}, deleted=${arb.deletedCount}} " +
+ s"serve=${arb.serve.size} stalePolicy=${if (staleFailPolicy) "fail" else "fallback"}")
+ if (staleFailPolicy && !arb.stale.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search rejected ${arb.staleCount} stale finalist(s) because " +
+ "vector.stale.locator.policy=FAIL.")
+ }
+ arb.serve.asScala.toSeq ++ arb.stale.asScala.toSeq
+ } else {
+ topCandidateList
+ }
+
+ if (serveCandidates.isEmpty) {
+ emptyDf
+ } else {
+ val candidatesByPartition = serveCandidates
+ .groupBy(candidate => Option(candidate.getLocation).map(_.getPartitionPath).getOrElse(""))
+ .map { case (partitionPath, candidates) =>
+ partitionPath -> candidates.flatMap(candidate => Option(candidate.getLocation).map(_.getFileId)).toSet
+ }
+ val fileSliceMap = preResolveFileSlices(
+ ctx.metadataTable, ctx.metaClient, ctx.timeline, ctx.latestInstantTime, candidatesByPartition)
+ if (fileSliceMap.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search resolved zero file slices for ${serveCandidates.size} rerank candidates " +
+ s"across ${candidatesByPartition.values.map(_.size).sum} file groups. Candidate locators are present, so this indicates " +
+ s"a file-slice resolution bug or stale vector index.")
+ }
+ val bFileSliceMap = spark.sparkContext.broadcast(fileSliceMap)
+ // MOR log-resident candidates carry their raw embedding in a log block, not the base Parquet
+ // file, so the positional fetcher cannot read them. Split them out and materialize them via a
+ // merged (base + log) read; base-resident candidates keep the fast positional path.
+ val (logResidentCandidates, baseCandidates) = serveCandidates.partition { candidate =>
+ bFileSliceMap.value.get(candidate.getLocation.getFileId).exists { slice =>
+ isLogResidentCandidate(candidate, slice.getBaseInstantTime, slice.getLogFiles.findAny().isPresent)
+ }
+ }
+ val exactFetchParallelism = math.max(
+ 1,
+ math.min(refineK, math.max(1, spark.sparkContext.defaultParallelism * 4)))
+ val groupedCandidates = baseCandidates.groupBy(_.getLocation.getFileId).toSeq.sortBy(_._1)
+ val targetBatchesPerGroup = math.max(1, exactFetchParallelism / math.max(1, groupedCandidates.size))
+ val fetchBatches = groupedCandidates.flatMap { case (fgId, candidates) =>
+ val sorted = candidates.sortBy(candidate => candidate.getLocation.getPosition)
+ val batchSize = math.max(1, math.ceil(sorted.size.toDouble / targetBatchesPerGroup).toInt)
+ sorted.grouped(batchSize).map(batch => (fgId, batch)).toSeq
+ }
+ val byFetchBatch: RDD[(String, Seq[ScoredVectorPostingMatch])] =
+ if (fetchBatches.isEmpty) {
+ spark.sparkContext.emptyRDD[(String, Seq[ScoredVectorPostingMatch])]
+ } else {
+ spark.sparkContext.parallelize(fetchBatches, math.min(exactFetchParallelism, fetchBatches.size))
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read_plan] rerankCandidates=${serveCandidates.size} " +
+ s"candidateGroups=${groupedCandidates.size} resolvedSlices=${fileSliceMap.size} " +
+ s"fileGroupPartitions=${byFetchBatch.getNumPartitions} exactFetchParallelism=$exactFetchParallelism")
+
+ val vectorOnlyProjection = mergedOptions.get(VECTOR_ONLY_PROJECTION_OPT)
+ .orElse(spark.conf.getOption(VECTOR_ONLY_PROJECTION_SPARK_CONF))
+ .exists(_.equalsIgnoreCase("true"))
+ val fetchSkippedLogResidentAcc: LongAccumulator =
+ spark.sparkContext.longAccumulator("vector_fetch_skipped_log_resident")
+ val fetchKeyMismatchAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_fetch_key_mismatch")
+ val rowsMaterializedAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_rows_materialized")
+ val exactRows: RDD[InternalRow] = byFetchBatch.mapPartitions { batches =>
+ val startNs = System.nanoTime()
+ val rows = ArrayBuffer.empty[InternalRow]
+ var batchCount = 0
+ var resolvedSlices = 0
+ var missingSlices = 0
+ var candidateKeys = 0L
+ var baseFiles = 0
+ var logFiles = 0L
+ var totalSliceBytes = 0L
+ var staleCandidates = 0L
+ var fetchSkippedLogResident = 0L
+ var fetchKeyMismatch = 0L
+ var keyLookupRowsDecoded = 0L
+ var metrics = PositionalParquetFetcher.FetchMetrics()
+ val fetcher = new PositionalParquetFetcher(
+ storageConf.newInstance().unwrap(),
+ embeddingCol,
+ vectorSchema,
+ exactQueryVector,
+ ctx.metric,
+ outputSchema,
+ mergedOptions.getOrElse("vector.exact_fetch.offset_index_missing_policy", "fallback").equalsIgnoreCase("fail"),
+ vectorOnlyProjection,
+ VectorIndexOptions.shouldVerifyFetchKeys(mergedOptions.asJava))
+ val keyLocator = new ParquetRecordKeyLocator(storageConf.newInstance().unwrap())
+ while (batches.hasNext) {
+ val (fgId, candidates) = batches.next()
+ batchCount += 1
+ val fileSliceOpt = bFileSliceMap.value.get(fgId)
+ if (fileSliceOpt.isEmpty) {
+ missingSlices += 1
+ throw new HoodieAnalysisException(
+ s"Vector exact search could not resolve file group '$fgId' for ${candidates.size} rerank candidates. " +
+ s"No record-level-index fallback is active in this path.")
+ } else {
+ resolvedSlices += 1
+ val fileSlice = fileSliceOpt.get
+ val logFileCount = fileSlice.getLogFiles.count()
+ val currentBaseInstantTime = fileSlice.getBaseInstantTime
+ val logResident = candidates.filter(candidate =>
+ isLogResidentCandidate(candidate, currentBaseInstantTime, logFileCount > 0))
+ if (logResident.nonEmpty) {
+ fetchSkippedLogResident += logResident.size
+ fetchSkippedLogResidentAcc.add(logResident.size)
+ LOG.warn(
+ s"Vector exact fetch skipped ${logResident.size} log-resident candidate(s) in file group '$fgId'; " +
+ "they remain a result shortfall and the retained candidate pool will refill where candidates remain")
+ }
+ val liveCandidates = candidates.filterNot(logResident.toSet)
+ staleCandidates += liveCandidates.count(candidate =>
+ candidate.getArbiterDecision == VectorIndexArbiter.Decision.STALE)
+ if (liveCandidates.nonEmpty) {
+ val positionalCandidates = liveCandidates.filter { candidate =>
+ val location = candidate.getLocation
+ location != null && location.getPosition >= 0 &&
+ candidate.getArbiterDecision != VectorIndexArbiter.Decision.STALE &&
+ (location.getInstantTime == null || location.getInstantTime == currentBaseInstantTime)
+ }
+ val keyCandidates = liveCandidates.filterNot(positionalCandidates.toSet)
+ val recordKeys = liveCandidates.map(_.getRecordKey).toSet
+ candidateKeys += recordKeys.size
+ if (fileSlice.getBaseFile.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact key fallback requires a base Parquet file for file group '$fgId'; " +
+ "the remaining candidates are not classified as log-resident.")
+ }
+ baseFiles += 1
+ logFiles += logFileCount
+ totalSliceBytes += fileSlice.getTotalFileSize
+ val baseFilePath = fileSlice.getBaseFile.get.getPath
+ val locatedByKey = keyLocator.locate(baseFilePath, keyCandidates)
+ keyLookupRowsDecoded += locatedByKey.metrics.rowsDecoded
+ if (locatedByKey.metrics.located != keyCandidates.size) {
+ throw new HoodieAnalysisException(
+ s"Vector exact key fallback located ${locatedByKey.metrics.located} of ${keyCandidates.size} " +
+ s"candidate key(s) in the live base file for file group '$fgId'.")
+ }
+ val firstFetch = fetcher.fetch(baseFilePath, positionalCandidates ++ locatedByKey.candidates)
+ rows ++= firstFetch.rows
+ metrics = metrics.add(firstFetch.metrics)
+ rowsMaterializedAcc.add(firstFetch.metrics.rowsMaterialized)
+ if (firstFetch.retryByKey.nonEmpty) {
+ fetchKeyMismatch += firstFetch.retryByKey.size
+ fetchKeyMismatchAcc.add(firstFetch.retryByKey.size)
+ val retries = keyLocator.locate(baseFilePath, firstFetch.retryByKey)
+ keyLookupRowsDecoded += retries.metrics.rowsDecoded
+ if (retries.metrics.located != firstFetch.retryByKey.size) {
+ throw new HoodieAnalysisException(
+ s"Vector exact fetch could not relocate ${firstFetch.retryByKey.size - retries.metrics.located} " +
+ s"key mismatch(es) in file group '$fgId'.")
+ }
+ val retryFetch = fetcher.fetch(baseFilePath, retries.candidates)
+ if (retryFetch.retryByKey.nonEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact fetch could not self-heal ${retryFetch.retryByKey.size} key mismatch(es) " +
+ s"in file group '$fgId' after key-based relocation.")
+ }
+ rows ++= retryFetch.rows
+ metrics = metrics.add(retryFetch.metrics)
+ rowsMaterializedAcc.add(retryFetch.metrics.rowsMaterialized)
+ }
+ }
+ }
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read] rerankCandidates=$candidateKeys candidateFileCount=$resolvedSlices " +
+ s"fetchBatches=$batchCount resolvedSlices=$resolvedSlices missingSlices=$missingSlices " +
+ s"rowsDecoded=${metrics.rowsDecoded} rowsMaterialized=${metrics.rowsMaterialized} " +
+ s"keyLookupRowsDecoded=$keyLookupRowsDecoded fetchKeyMismatch=$fetchKeyMismatch " +
+ s"fetchSkippedLogResident=$fetchSkippedLogResident " +
+ s"rowGroupsTotal=${metrics.rowGroupsTotal} rowGroupsSelected=${metrics.rowGroupsSelected} " +
+ s"pagesInSelectedRowGroups=${metrics.pagesInSelectedRowGroups} pagesSelected=${metrics.pagesSelected} " +
+ s"offsetIndexHits=${metrics.offsetIndexHits} offsetIndexMissing=${metrics.offsetIndexMissing} " +
+ s"rangedGets=${metrics.rangedGets} pageBytesFetched=${metrics.pageBytesFetched} " +
+ s"physicalBytesRead=${metrics.physicalBytesRead} footerReads=${metrics.footerReads} footerCacheHits=${metrics.footerCacheHits} " +
+ s"fetchKeyMismatchMetric=${metrics.fetchKeyMismatch} staleCandidates=$staleCandidates " +
+ s"baseFiles=$baseFiles logFiles=$logFiles " +
+ s"sliceBytes=$totalSliceBytes footerMs=${metrics.footerMs} offsetIndexMs=${metrics.offsetIndexMs} " +
+ s"fetchWaitMs=${metrics.fetchWaitMs} decodeMs=${metrics.decodeMs} scoreMs=${metrics.scoreMs} elapsedMs=${elapsedMs(startNs)}")
+ rows.iterator
+ }
+ val logResidentRows: Seq[InternalRow] =
+ if (logResidentCandidates.isEmpty) {
+ Seq.empty
+ } else {
+ val dataSchema = new TableSchemaResolver(ctx.metaClient).getTableSchema(true)
+ val scorer = new VectorExactScorer(vectorSchema, exactQueryVector, ctx.metric)
+ val logFetcher = new LogResidentVectorFetcher(
+ storageConf, ctx.metaClient, dataSchema, ctx.latestInstantTime,
+ embeddingCol, outputSchema, scorer)
+ val fetched = logResidentCandidates
+ .groupBy(candidate =>
+ (Option(candidate.getLocation.getPartitionPath).getOrElse(""), candidate.getLocation.getFileId))
+ .flatMap { case ((partitionPath, fgId), candidates) =>
+ logFetcher.fetchSlice(partitionPath, fileSliceMap(fgId), candidates.map(_.getRecordKey).toSet)
Review Comment:
**Fixed** `10f50b3319`. Merged (base+log) reads run in executor tasks over the broadcast slice map, not on the driver.
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorSearchAlgorithms.scala:
##########
@@ -0,0 +1,413 @@
+/*
+ * 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.spark.sql.hudi.analysis
+
+import org.apache.hudi.SparkAdapterSupport
Review Comment:
**Fixed** `3cde5283ea`. Removed.
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorSearchAlgorithms.scala:
##########
@@ -0,0 +1,413 @@
+/*
+ * 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.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.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, 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}
+
+case class VectorSearchTable(df: DataFrame, basePath: Option[String])
+
+/**
+ * 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]]
+ *
+ * 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()`.
+ */
+trait VectorSearchAlgorithm {
+
+ /** Human-readable name for error messages and logging. */
+ def name: String
+
+ /**
+ * Build a plan that finds the k nearest corpus rows to a single query vector.
+ *
+ * @param spark active SparkSession
+ * @param corpusTable resolved corpus table (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)
+ * @return an analyzed LogicalPlan whose output matches the single-query schema contract
+ */
+ def buildSingleQueryPlan(
+ spark: SparkSession,
+ corpusTable: VectorSearchTable,
+ embeddingCol: String,
+ queryVector: Array[Double],
+ k: Int,
+ metric: DistanceMetric.Value,
+ runtimeOptions: Map[String, String] = Map.empty): LogicalPlan
+
+ /**
+ * Build a plan that finds the k nearest corpus rows for each row in the query table.
+ *
+ * @param spark active SparkSession
+ * @param corpusTable resolved corpus table
+ * @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)
+ * @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,
+ corpusTable: VectorSearchTable,
+ corpusEmbeddingCol: String,
+ queryDf: DataFrame,
+ queryEmbeddingCol: String,
+ k: Int,
+ metric: DistanceMetric.Value,
+ runtimeOptions: Map[String, String] = Map.empty): LogicalPlan
+}
+
+/**
+ * Resolves [[SearchAlgorithm]] values to [[VectorSearchAlgorithm]] implementations
+ * and provides shared validation helpers used across algorithms.
+ */
+object HoodieVectorSearchPlanBuilder {
+
+ private[analysis] val LOG = LoggerFactory.getLogger(getClass)
+
+ val DISTANCE_COL = "_hudi_distance"
+ val FILE_GROUP_ID_COL = "_hudi_file_group_id"
+ val VECTOR_ONLY_PROJECTION_OPT = "vector.exact_fetch.vector_only_projection"
+ val VECTOR_ONLY_PROJECTION_SPARK_CONF = "spark.hudi.vector.exact_fetch.vector_only_projection"
+ 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_"
+
+ private[analysis] def elapsedMs(startNs: Long): Long = (System.nanoTime() - startNs) / 1000000L
+
+ private[analysis] def summarizeInts(values: Seq[Int], limit: Int = 16): String = {
+ val preview = values.take(limit).mkString("[", ", ", "]")
+ if (values.length > limit) s"$preview...(${values.length} total)" else preview
+ }
+
+ private[analysis] def summarizeShardCounts(shardCounts: collection.Map[Int, Int], limit: Int = 16): String = {
+ val preview = shardCounts.toSeq.sortBy(_._1).take(limit).map { case (clusterId, count) =>
+ s"$clusterId->$count"
+ }.mkString("[", ", ", "]")
+ if (shardCounts.size > limit) s"$preview...(${shardCounts.size} total)" else preview
+ }
+
+ private[analysis] def logRddLineage(label: String, rdd: RDD[_]): Unit = {
+ Try(rdd.toDebugString) match {
+ case Success(debug) =>
+ LOG.info(s"[vector_search][$label] lineage:\n$debug")
+ case Failure(error) =>
+ LOG.warn(s"[vector_search][$label] unable to fetch RDD lineage", error)
+ }
+ }
+
+ /** Resolve a [[SearchAlgorithm]] enum value to its implementation. */
+ def resolveAlgorithm(algorithm: SearchAlgorithm.Value): VectorSearchAlgorithm = algorithm match {
+ case SearchAlgorithm.BRUTE_FORCE => BruteForceSearchAlgorithm
+ case SearchAlgorithm.IVF_RABITQ_MDT => IvfRaBitQMdtSearchAlgorithm
+ case other => throw new HoodieAnalysisException(
+ s"Unsupported search algorithm: $other")
+ }
+
+ 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>")
+ }
+ }
+
+ /**
+ * 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 = {
+ extractVectorSchema(df, embeddingCol).map(_.getDimension).foreach { corpusDim =>
+ if (corpusDim != queryDim) {
+ throw new HoodieAnalysisException(
+ s"Query vector dimension ($queryDim) does not match " +
+ s"corpus embedding dimension ($corpusDim) for column '$embeddingCol'")
+ }
+ }
+ }
+
+ /**
+ * 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) {
+ 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>).")
+ }
+ }
+
+ /**
+ * 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 = {
+ (extractVectorSchema(corpusDf, corpusCol).map(_.getDimension), extractVectorSchema(queryDf, queryCol).map(_.getDimension)) 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
+ }
+ }
+
+ 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 =>
+ throw new HoodieAnalysisException(
+ s"Embedding column '$colName' has type $other, expected an array type")
+ }
+ }
+
+ /** Extracts VECTOR schema from column metadata, if present. */
+ private[analysis] def extractVectorSchema(df: DataFrame, colName: String): Option[HoodieSchema.Vector] = {
+ 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)
+ 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[analysis] def quoteIdentifier(name: String): String = s"`${name.replace("`", "``")}`"
+}
+
+/**
+ * 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.
+ */
+object BruteForceSearchAlgorithm extends VectorSearchAlgorithm {
+
+ import HoodieVectorSearchPlanBuilder._
+
+ override val name: String = "brute_force"
+
+ 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 def filterOption(runtimeOptions: Map[String, String]): Option[String] =
+ runtimeOptions.get(HoodieVectorSearchTableValuedFunction.BRUTE_FORCE_FILTER_OPT)
+
+ private def maxDistanceOption(runtimeOptions: Map[String, String]): Option[Double] =
+ runtimeOptions.get(HoodieVectorSearchTableValuedFunction.BRUTE_FORCE_MAX_DISTANCE_OPT).map(_.toDouble)
+
+ override def buildSingleQueryPlan(
+ spark: SparkSession,
+ corpusTable: VectorSearchTable,
+ embeddingCol: String,
+ queryVector: Array[Double],
+ k: Int,
+ metric: DistanceMetric.Value,
+ runtimeOptions: Map[String, String] = Map.empty): LogicalPlan = {
+ val corpusDf = corpusTable.df
+ validateEmbeddingColumn(corpusDf, embeddingCol)
+ validateQueryVectorDimension(corpusDf, embeddingCol, queryVector.length)
+
+ val elemType = getElementType(corpusDf, embeddingCol)
+ var filteredDf = corpusDf.filter(col(embeddingCol).isNotNull)
Review Comment:
**Fixed** `3cde5283ea`. Replaced with two small helpers.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4384291517
##########
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.
Compatibility follows the existing MDT dispatch contract: pre-vector binaries reject the unknown vector
partition in `MetadataPartitionType.fromPartitionPath`; tables with a vector index are not writable by
those binaries. Same-version writers maintain it unconditionally, including clients initialized before
CREATE INDEX (covered by the production-client regression). Per-commit markers are audit records;
strict validation is read-only, opt-in, defaults off, and has no correctness role. Its policy defaults
WARN in both modes. There is no per-commit frontier/barrier or catch-up protocol. Diagnostic gaps are
repaired by dropping and recreating the index, not an unimplemented replay operation.
This describes the complete maintenance stack; base alone still has a buildUpdate stub.
##########
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.
Latest Spark 4.0 vector-only validation: the inflight-bootstrap case first raised `TableNotFoundException` while constructing the commit's table client, then passed on retry. This concurrent config-read issue is retained as a separate investigation, not claimed fixed by the vector patch. The same-version pre-index client and four lifecycle/rollback cases passed without reruns. The stacked PR records the validation details.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "yihua (via GitHub)" <gi...@apache.org>.
yihua commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4163156384
##########
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:
See https://github.com/apache/hudi/pull/19802#issuecomment-5946183886
Wrt partitioned RLI, `readRecordIndexLocationsWithKeys(keys, partition)` has an implicit precondition that all keys in a batch hash to the same RLI shard, and every current caller (Spark write/read, Flink read) groups the lookup keys by partition AND shard before calling. The vector lookup caller here does not follow this contract - it invokes the method after grouping the keys by data table partition only, without shard ID within the partition, which breaks on partitioned RLI with more than one shard per data table partition.
The fix is to group finalist keys by mapRecordKeyToFileGroupIndex within each partition, the way the Spark/Flink read path does.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4384296622
##########
hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataPayload.java:
##########
@@ -253,6 +300,410 @@ protected HoodieMetadataPayload(String key, HoodieSecondaryIndexInfo secondaryIn
this(key, MetadataPartitionType.SECONDARY_INDEX.getRecordType(), null, null, null, null, secondaryIndexMetadata, secondaryIndexMetadata.getIsDeleted());
}
+ protected HoodieMetadataPayload(String key, Object vectorIndexInfo) {
+ this.key = key;
+ this.type = MetadataPartitionType.VECTOR_INDEX.getRecordType();
+ this.vectorIndexMetadata = vectorIndexInfo;
+ this.isDeletedRecord = vectorIndexInfo instanceof HoodieVectorIndexTombstone;
+ }
+
+ /**
+ * Create the singleton reader-visible generation pointer.
+ */
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexActiveManifestRecord(
+ Integer activeGeneration, String metadataPartitionPath) {
+ String recordKey = VectorIndexMetadataKey.activeManifest();
+ HoodieVectorIndexActiveManifest manifest = new HoodieVectorIndexActiveManifest(1, activeGeneration);
+ return new HoodieAvroRecord<>(
+ new HoodieKey(recordKey, metadataPartitionPath),
+ new HoodieMetadataPayload(recordKey, manifest));
+ }
+
+ /**
+ * Create the generation-one centroid record for the given index partition.
+ */
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexCentroidsRecord(
+ ByteBuffer centroidBytes, String partitionPath) {
+ HoodieVectorIndexCentroids centroids = new HoodieVectorIndexCentroids(
+ ByteBuffer.allocate(0),
+ centroidBytes,
+ ByteBuffer.allocate(0));
+ String recordKey = VectorIndexMetadataKey.centroids(1, 0);
+ HoodieMetadataPayload payload = new HoodieMetadataPayload(recordKey, centroids);
+ HoodieKey key = new HoodieKey(recordKey, partitionPath);
+ return new HoodieAvroRecord<>(key, payload);
+ }
+
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexCentroidsRecord(
+ int generation,
+ int chunk,
+ ByteBuffer clusterIds,
+ ByteBuffer centroidBytes,
+ ByteBuffer clusterRadii,
+ String partitionPath) {
+ String recordKey = VectorIndexMetadataKey.centroids(generation, chunk);
+ HoodieVectorIndexCentroids centroids = new HoodieVectorIndexCentroids(
+ clusterIds, centroidBytes, clusterRadii);
+ return new HoodieAvroRecord<>(
+ new HoodieKey(recordKey, partitionPath),
+ new HoodieMetadataPayload(recordKey, centroids));
+ }
+
+ /**
+ * Create the generation-one quantizer metadata record for the given index partition.
+ */
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexQuantizerMetadataRecord(
+ String quantizerType,
+ int quantizedCodeBytes,
+ long randomSeed,
+ boolean assumeNormalized,
+ String partitionPath) {
+ return createVectorIndexQuantizerMetadataRecord(
+ quantizerType,
+ quantizedCodeBytes,
+ 1,
+ randomSeed,
+ assumeNormalized,
+ partitionPath);
+ }
+
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexQuantizerMetadataRecord(
+ String quantizerType,
+ int quantizedCodeBytes,
+ int rabitqBits,
+ long randomSeed,
+ boolean assumeNormalized,
+ String partitionPath) {
+ return createVectorIndexQuantizerMetadataRecord(1, 0, quantizerType, randomSeed, null, partitionPath);
Review Comment:
**Fixed** `3cde5283ea`. The overloads that discarded arguments are removed.
##########
hudi-common/src/main/java/org/apache/hudi/metadata/HoodieTableMetadataUtil.java:
##########
@@ -213,6 +215,71 @@ public class HoodieTableMetadataUtil {
HoodieRecord.HoodieMetadataField.PARTITION_PATH_METADATA_FIELD.getFieldName(),
HoodieRecord.HoodieMetadataField.COMMIT_TIME_METADATA_FIELD.getFieldName()));
+ public static String getVectorIndexGenerationManifestKey(int generationId) {
+ return VectorIndexMetadataKey.manifest(generationId);
+ }
+
+ public static boolean isVectorIndexGenerationManifestKey(String recordKey) {
+ return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_MANIFEST);
+ }
+
+ public static boolean isVectorIndexQuantizerKey(String recordKey) {
+ return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_QUANTIZER);
+ }
+
+ public static boolean isVectorIndexCentroidsKey(String recordKey) {
+ return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_CENTROIDS);
+ }
+
+ public static String getVectorIndexClusterKey(int generationId, int clusterId) {
+ return VectorIndexMetadataKey.clusterStats(generationId, clusterId);
+ }
+
+ public static boolean isVectorIndexClusterKey(String recordKey) {
+ return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_CLUSTER_STATS);
+ }
+
+ public static String getVectorIndexPostingKey(int generationId, int clusterId, int shardId, String recordKey) {
+ return VectorIndexMetadataKey.postingDelta(generationId, clusterId, shardId, recordKey);
+ }
+
+ public static String getVectorIndexPostingPrefix(int generationId, int clusterId) {
+ return VectorIndexMetadataKey.postingPrefix(generationId, clusterId, 0).substring(0, 9);
Review Comment:
**Fixed** `c8b06883b1`. `VectorIndexMetadataKey.postingClusterPrefix` / `POSTING_CLUSTER_PREFIX_LENGTH`.
##########
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()));
+ }
+ return finalists.map(candidate ->
+ arbitrateCandidate(candidate, currentLocations.get(candidate.getRecordKey())));
+ }
+
+ private static ScoredVectorPostingMatch arbitrateCandidate(
+ ScoredVectorPostingMatch candidate,
+ HoodieRecordGlobalLocation current) {
+ VectorIndexArbiter.Decision decision = VectorIndexArbiter.classify(
+ candidate.getPartitionPath(),
+ candidate.getFileGroupId(),
+ candidate.getBaseInstantTime(),
+ current);
+ HoodieRecordGlobalLocation resolved;
+ switch (decision) {
+ case SERVE:
+ resolved = candidate.getLocation() != null ? candidate.getLocation() : current;
+ break;
+ case STALE:
+ resolved = current;
+ break;
+ case DELETED:
+ default:
+ resolved = null;
+ break;
+ }
+ return candidate.withArbiterVerdict(decision, resolved);
+ }
+
+ /**
+ * Driver-side finalist arbiter core: classify each already-materialized finalist against a
+ * pre-resolved map of current RLI locations (record key -> location, or absent for an RLI miss).
+ * Pure and Spark-free so it is directly unit-testable; the {@link HoodieTableMetadata} overload
+ * performs the batched RLI lookup and delegates here.
+ *
+ * <p>Implements the RFC-109 arbiter output contract (see {@link VectorIndexArbiter}):
+ * hit+match -> SERVE (positional trust preserved via the posting's own location when present),
+ * hit+differ -> STALE (resolved to the live RLI location), miss -> DELETED (dropped).
+ */
+ static VectorIndexMdtSearchUtils.ArbitrationResult arbitrateMaterializedFinalists(
+ List<ScoredVectorPostingMatch> finalists,
+ Map<String, HoodieRecordGlobalLocation> currentLocations) {
+ List<ScoredVectorPostingMatch> serve = new ArrayList<>();
+ List<ScoredVectorPostingMatch> stale = new ArrayList<>();
+ long deleted = 0L;
+ for (ScoredVectorPostingMatch candidate : finalists) {
+ HoodieRecordGlobalLocation current = currentLocations.get(candidate.getRecordKey());
+ VectorIndexArbiter.Decision decision = VectorIndexArbiter.classify(
+ candidate.getPartitionPath(),
+ candidate.getFileGroupId(),
+ candidate.getBaseInstantTime(),
+ current);
Review Comment:
**Fixed** `c8b06883b1`. The loop buckets `arbitrateCandidate`'s verdict; the rules are defined once.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "yihua (via GitHub)" <gi...@apache.org>.
yihua commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-6485892114
Hi @chrevanthreddy, thanks for updating the PR! I'll review the changes again, and let's get the initial implementation merged with a tight scope, with independent follow-ups, so it's easier for both reviewers and for you as the author.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "codecov-commenter (via GitHub)" <gi...@apache.org>.
codecov-commenter commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-5860340304
## [Codecov](https://app.codecov.io/gh/apache/hudi/pull/20107?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) Report
:x: Patch coverage is `0.45920%` with `2818 lines` in your changes missing coverage. Please review.
:white_check_mark: Project coverage is 78.45%. Comparing base ([`513edcf`](https://app.codecov.io/gh/apache/hudi/commit/513edcf8c5e0a811699d2d9db2645252cd158199?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)) to head ([`4204280`](https://app.codecov.io/gh/apache/hudi/commit/4204280759551a140914b8a393d0420206489542?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)).
:warning: Report is 10 commits behind head on master.
| [Files with missing lines](https://app.codecov.io/gh/apache/hudi/pull/20107?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Patch % | Lines |
|---|---|---|
| [...common/index/vector/VectorIndexMdtSearchUtils.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fcommon%2Findex%2Fvector%2FVectorIndexMdtSearchUtils.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvY29tbW9uL2luZGV4L3ZlY3Rvci9WZWN0b3JJbmRleE1kdFNlYXJjaFV0aWxzLmphdmE=) | 0.00% | [603 Missing :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...apache/hudi/common/index/vector/RaBitQEncoder.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fcommon%2Findex%2Fvector%2FRaBitQEncoder.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvY29tbW9uL2luZGV4L3ZlY3Rvci9SYUJpdFFFbmNvZGVyLmphdmE=) | 0.00% | [271 Missing :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [.../common/index/vector/VectorIndexMetadataCache.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fcommon%2Findex%2Fvector%2FVectorIndexMetadataCache.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvY29tbW9uL2luZGV4L3ZlY3Rvci9WZWN0b3JJbmRleE1ldGFkYXRhQ2FjaGUuamF2YQ==) | 0.00% | [224 Missing :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [.../hudi/common/index/vector/PostingBlockBuilder.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fcommon%2Findex%2Fvector%2FPostingBlockBuilder.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvY29tbW9uL2luZGV4L3ZlY3Rvci9Qb3N0aW5nQmxvY2tCdWlsZGVyLmphdmE=) | 0.00% | [149 Missing :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...rg/apache/hudi/metadata/HoodieMetadataPayload.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fmetadata%2FHoodieMetadataPayload.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvbWV0YWRhdGEvSG9vZGllTWV0YWRhdGFQYXlsb2FkLmphdmE=) | 0.86% | [114 Missing and 1 partial :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...e/hudi/common/index/vector/VectorIndexOptions.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fcommon%2Findex%2Fvector%2FVectorIndexOptions.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvY29tbW9uL2luZGV4L3ZlY3Rvci9WZWN0b3JJbmRleE9wdGlvbnMuamF2YQ==) | 0.00% | [111 Missing :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...che/hudi/common/index/vector/PostingBlockView.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fcommon%2Findex%2Fvector%2FPostingBlockView.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvY29tbW9uL2luZGV4L3ZlY3Rvci9Qb3N0aW5nQmxvY2tWaWV3LmphdmE=) | 0.00% | [104 Missing :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [.../common/index/vector/VectorIndexMetadataSerde.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fcommon%2Findex%2Fvector%2FVectorIndexMetadataSerde.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvY29tbW9uL2luZGV4L3ZlY3Rvci9WZWN0b3JJbmRleE1ldGFkYXRhU2VyZGUuamF2YQ==) | 0.00% | [98 Missing :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...e/hudi/common/index/vector/VectorPostingSerde.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fcommon%2Findex%2Fvector%2FVectorPostingSerde.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvY29tbW9uL2luZGV4L3ZlY3Rvci9WZWN0b3JQb3N0aW5nU2VyZGUuamF2YQ==) | 0.00% | [94 Missing :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| [...g/apache/hudi/metadata/VectorIndexMetadataKey.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fmetadata%2FVectorIndexMetadataKey.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvbWV0YWRhdGEvVmVjdG9ySW5kZXhNZXRhZGF0YUtleS5qYXZh) | 0.00% | [94 Missing :warning: ](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) |
| ... and [32 more](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | |
<details><summary>Additional details and impacted files</summary>
```diff
@@ Coverage Diff @@
## master #20107 +/- ##
============================================
- Coverage 80.37% 78.45% -1.92%
- Complexity 34794 35296 +502
============================================
Files 2545 2579 +34
Lines 142722 147842 +5120
Branches 17336 18225 +889
============================================
+ Hits 114708 115988 +1280
- Misses 20104 23826 +3722
- Partials 7910 8028 +118
```
| [Components](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=components&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [hudi-common](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `77.51% <0.47%> (-6.40%)` | :arrow_down: |
| [hudi-client](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `81.74% <0.00%> (-1.69%)` | :arrow_down: |
| [hudi-flink](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `85.68% <ø> (ø)` | |
| [hudi-spark-datasource](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `73.76% <ø> (-0.03%)` | :arrow_down: |
| [hudi-utilities](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `78.15% <ø> (ø)` | |
| [hudi-cli](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `69.99% <ø> (ø)` | |
| [hudi-hadoop](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `70.98% <ø> (ø)` | |
| [hudi-sync](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `76.02% <ø> (ø)` | |
| [hudi-io](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `81.26% <ø> (-0.32%)` | :arrow_down: |
| [hudi-timeline-service](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `82.67% <ø> (-0.68%)` | :arrow_down: |
| [hudi-cloud](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `81.00% <ø> (ø)` | |
| [hudi-kafka-connect](https://app.codecov.io/gh/apache/hudi/pull/20107/components?src=pr&el=component&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `53.96% <ø> (ø)` | |
| [Flag](https://app.codecov.io/gh/apache/hudi/pull/20107/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [hadoop-mr-java-client](https://app.codecov.io/gh/apache/hudi/pull/20107/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | `42.27% <0.45%> (-1.58%)` | :arrow_down: |
Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#carryforward-flags-in-the-pull-request-comment) to find out more.
| [Files with missing lines](https://app.codecov.io/gh/apache/hudi/pull/20107?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | Coverage Δ | |
|---|---|---|
| [.../org/apache/hudi/index/HoodieSparkIndexClient.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-client%2Fhudi-spark-client%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Findex%2FHoodieSparkIndexClient.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jbGllbnQvaHVkaS1zcGFyay1jbGllbnQvc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvaW5kZXgvSG9vZGllU3BhcmtJbmRleENsaWVudC5qYXZh) | `84.32% <ø> (ø)` | |
| [...pache/hudi/metadata/index/SparkIndexerSupport.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-client%2Fhudi-spark-client%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fmetadata%2Findex%2FSparkIndexerSupport.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jbGllbnQvaHVkaS1zcGFyay1jbGllbnQvc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvbWV0YWRhdGEvaW5kZXgvU3BhcmtJbmRleGVyU3VwcG9ydC5qYXZh) | `76.47% <ø> (ø)` | |
| [...ogical/HoodieVectorSearchTableValuedFunction.scala](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-spark-datasource%2Fhudi-spark-common%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fspark%2Fsql%2Fcatalyst%2Fplans%2Flogical%2FHoodieVectorSearchTableValuedFunction.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1zcGFyay1kYXRhc291cmNlL2h1ZGktc3BhcmstY29tbW9uL3NyYy9tYWluL3NjYWxhL29yZy9hcGFjaGUvc3Bhcmsvc3FsL2NhdGFseXN0L3BsYW5zL2xvZ2ljYWwvSG9vZGllVmVjdG9yU2VhcmNoVGFibGVWYWx1ZWRGdW5jdGlvbi5zY2FsYQ==) | `69.04% <ø> (-1.79%)` | :arrow_down: |
| [...rk/sql/hudi/analysis/HoodieSparkBaseAnalysis.scala](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-spark-datasource%2Fhudi-spark-common%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fspark%2Fsql%2Fhudi%2Fanalysis%2FHoodieSparkBaseAnalysis.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1zcGFyay1kYXRhc291cmNlL2h1ZGktc3BhcmstY29tbW9uL3NyYy9tYWluL3NjYWxhL29yZy9hcGFjaGUvc3Bhcmsvc3FsL2h1ZGkvYW5hbHlzaXMvSG9vZGllU3BhcmtCYXNlQW5hbHlzaXMuc2NhbGE=) | `69.86% <ø> (-1.27%)` | :arrow_down: |
| [.../hudi/analysis/HoodieVectorSearchPlanBuilder.scala](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-spark-datasource%2Fhudi-spark-common%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fspark%2Fsql%2Fhudi%2Fanalysis%2FHoodieVectorSearchPlanBuilder.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1zcGFyay1kYXRhc291cmNlL2h1ZGktc3BhcmstY29tbW9uL3NyYy9tYWluL3NjYWxhL29yZy9hcGFjaGUvc3Bhcmsvc3FsL2h1ZGkvYW5hbHlzaXMvSG9vZGllVmVjdG9yU2VhcmNoUGxhbkJ1aWxkZXIuc2NhbGE=) | `75.00% <ø> (-0.64%)` | :arrow_down: |
| [.../apache/spark/sql/hudi/command/IndexCommands.scala](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-spark-datasource%2Fhudi-spark%2Fsrc%2Fmain%2Fscala%2Forg%2Fapache%2Fspark%2Fsql%2Fhudi%2Fcommand%2FIndexCommands.scala&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1zcGFyay1kYXRhc291cmNlL2h1ZGktc3Bhcmsvc3JjL21haW4vc2NhbGEvb3JnL2FwYWNoZS9zcGFyay9zcWwvaHVkaS9jb21tYW5kL0luZGV4Q29tbWFuZHMuc2NhbGE=) | `87.95% <ø> (-0.15%)` | :arrow_down: |
| [...org/apache/hudi/metadata/index/IndexerFactory.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-client%2Fhudi-client-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fmetadata%2Findex%2FIndexerFactory.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jbGllbnQvaHVkaS1jbGllbnQtY29tbW9uL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9odWRpL21ldGFkYXRhL2luZGV4L0luZGV4ZXJGYWN0b3J5LmphdmE=) | `84.00% <0.00%> (-3.50%)` | :arrow_down: |
| [...etadata/index/UnsupportedEngineIndexerSupport.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-client%2Fhudi-client-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fmetadata%2Findex%2FUnsupportedEngineIndexerSupport.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jbGllbnQvaHVkaS1jbGllbnQtY29tbW9uL3NyYy9tYWluL2phdmEvb3JnL2FwYWNoZS9odWRpL21ldGFkYXRhL2luZGV4L1Vuc3VwcG9ydGVkRW5naW5lSW5kZXhlclN1cHBvcnQuamF2YQ==) | `27.27% <0.00%> (-2.73%)` | :arrow_down: |
| [...a/org/apache/hudi/metadata/HoodieIndexVersion.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fmetadata%2FHoodieIndexVersion.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvbWV0YWRhdGEvSG9vZGllSW5kZXhWZXJzaW9uLmphdmE=) | `85.71% <0.00%> (-2.53%)` | :arrow_down: |
| [.../org/apache/hudi/metadata/VectorClusterRawKey.java](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree&filepath=hudi-common%2Fsrc%2Fmain%2Fjava%2Forg%2Fapache%2Fhudi%2Fmetadata%2FVectorClusterRawKey.java&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache#diff-aHVkaS1jb21tb24vc3JjL21haW4vamF2YS9vcmcvYXBhY2hlL2h1ZGkvbWV0YWRhdGEvVmVjdG9yQ2x1c3RlclJhd0tleS5qYXZh) | `0.00% <0.00%> (ø)` | |
| ... and [38 more](https://app.codecov.io/gh/apache/hudi/pull/20107?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache) | |
... and [19 files with indirect coverage changes](https://app.codecov.io/gh/apache/hudi/pull/20107/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=apache)
</details>
<details><summary> :rocket: New features to boost your workflow: </summary>
- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
- :package: [JS Bundle Analysis](https://docs.codecov.com/docs/javascript-bundle-analysis): Save yourself from yourself by tracking and limiting bundle sizes in JS merges.
</details>
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "yihua (via GitHub)" <gi...@apache.org>.
yihua commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4213559264
##########
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:
The `buildUpdate` and `buildClean` are not implemented yet, so the synchronous updates to the vector index do not work in this PR, in contrast to #19802.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "hudi-agent (via GitHub)" <gi...@apache.org>.
hudi-agent commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4236046529
##########
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:
🤖 non-blocking: the signature doc above advertises `algorithm_arg | NULL`, but for `ivf_rabitq_mdt` a `NULL` here goes through `requireStringLiteral` and is rejected (same at line 263 for the batch TVF); could you route it through `parseOptionalString` so an empty/NULL option list is accepted for both algorithms?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
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:
🤖 non-blocking: `parseTableIdentifier` rejects the `catalog.db.table` form the doc comment above promises, and the catch-all turns that into `None`, so `ivf_rabitq_mdt` then fails with the misleading "not a temporary view" message; could you take the location from the analyzed plan instead (e.g. `SparkAdapter.resolveHoodieTable(tableDf.queryExecution.analyzed)` + `HoodieSqlCommonUtils.getTableLocation`)?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
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:
🤖 non-blocking: only the structural keys are rejected here, so a mistyped runtime option (e.g. the retired `vector.query.nprobe=8` that CREATE INDEX rejects) is silently ignored at query time; would it make sense to validate the runtime keys against `VectorIndexOptions` too?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
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:
🤖 non-blocking: `VectorIndexOptions.isFinalistArbiterEnabled` is hardcoded to `true`, so this isn't dormant — approximate mode always collects the candidate keys to the driver for an RLI lookup and therefore always depends on the record index; could you align the comment with the option, or make it a real option?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
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:
🤖 non-blocking: building on @vinishjail97's point that the bootstrap currently records no frontier, comparing requested times here still misses a writer that was in flight when the bootstrap read the table (writer A requests a deltacommit at t3 and is still inflight when the bootstrap reads at t6 with last completed write t4; A completes at t7, but t3 < t4 so it counts as covered and the exact path silently serves a stale index under `fail`); would it make sense to compare completion times instead, e.g. `findInstantsModifiedAfterByCompletionTime` against the bootstrap's read instant?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorSearchAlgorithms.scala:
##########
@@ -0,0 +1,413 @@
+/*
+ * 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.spark.sql.hudi.analysis
+
+import org.apache.hudi.SparkAdapterSupport
Review Comment:
🤖 nit: it looks like most of these imports (SparkAdapterSupport, HoodieSparkEngineContext, TypedProperties, FileSlice, HoodieTableMetaClient, timeline/view, HoodieJavaRDD, JavaSparkContext, InternalRow, LongAccumulator, ArrayBuffer, StructType, ...) aren't used in this file anymore — could you prune them so a reader doesn't go looking for MDT/file-system-view code that lives in the IVF algorithm file?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
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:
🤖 non-blocking: `metadataCaches.synchronized` locks whatever map the var currently points at while the block reassigns it, so two concurrent queries can both pass the guard and load the same cache twice (harmless today, but the lock isn't doing what it looks like); a dedicated lock object or a `ConcurrentHashMap` with `compute` would make it real.
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexMdtSearchUtils.java:
##########
@@ -0,0 +1,1310 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.common.index.vector;
+
+import org.apache.hudi.avro.model.HoodieVectorIndexPostingBlock;
+import org.apache.hudi.common.data.HoodieData;
+import org.apache.hudi.common.data.HoodieListData;
+import org.apache.hudi.common.data.HoodiePairData;
+import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.model.HoodieRecordGlobalLocation;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.metadata.HoodieMetadataPayload;
+import org.apache.hudi.metadata.HoodieTableMetadata;
+import org.apache.hudi.metadata.RawKey;
+import org.apache.hudi.metadata.VectorIndexMetadataKey;
+import org.apache.hudi.metadata.VectorPostingPrefixRawKey;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.management.GarbageCollectorMXBean;
+import java.lang.management.ManagementFactory;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+
+/**
+ * Shared helper for MDT-native vector posting lookup and approximate candidate reduction.
+ *
+ * <p>This provides the missing second-stage building block after coarse IVF cluster pruning:
+ * exact HFile lookups of cluster metadata, prefix scans over {@code P|gen|cluster|shard|},
+ * approximate RaBitQ scoring, and direct file-slice targeting from posting payload metadata.
+ */
+public final class VectorIndexMdtSearchUtils {
+
+ private static final Logger LOG = LoggerFactory.getLogger(VectorIndexMdtSearchUtils.class);
+ private static final int TOP_K_REDUCER_KEY = 0;
+ private static final Comparator<ScoredVectorPostingMatch> BEST_FIRST = Comparator
+ .comparingDouble(ScoredVectorPostingMatch::getApproxDistance)
+ .thenComparing(ScoredVectorPostingMatch::getRecordKey);
+
+ private VectorIndexMdtSearchUtils() {
+ }
+
+ public static Map<Integer, Integer> readClusterShardCounts(HoodieTableMetadata metadataTable,
+ String indexPartition,
+ int generationId,
+ Collection<Integer> clusterIds) {
+ if (clusterIds == null || clusterIds.isEmpty()) {
+ return Collections.emptyMap();
+ }
+
+ Map<Integer, Integer> shardCounts = new HashMap<>();
+ for (Integer clusterId : clusterIds) {
+ shardCounts.put(clusterId, 1);
+ }
+ return shardCounts;
+ }
+
+ public static Map<Integer, Set<String>> readClusterFileGroups(HoodieTableMetadata metadataTable,
Review Comment:
🤖 nit: `readClusterFileGroups` always returns an empty map and `readClusterShardCounts` ignores the metadata table and hard-codes 1 shard per cluster; could you either drop these stubs or add a comment saying they are placeholders, so callers don't assume they actually read MDT?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorSearchAlgorithms.scala:
##########
@@ -0,0 +1,413 @@
+/*
+ * 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.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.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, 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}
+
+case class VectorSearchTable(df: DataFrame, basePath: Option[String])
+
+/**
+ * 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]]
+ *
+ * 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()`.
+ */
+trait VectorSearchAlgorithm {
+
+ /** Human-readable name for error messages and logging. */
+ def name: String
+
+ /**
+ * Build a plan that finds the k nearest corpus rows to a single query vector.
+ *
+ * @param spark active SparkSession
+ * @param corpusTable resolved corpus table (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)
+ * @return an analyzed LogicalPlan whose output matches the single-query schema contract
+ */
+ def buildSingleQueryPlan(
+ spark: SparkSession,
+ corpusTable: VectorSearchTable,
+ embeddingCol: String,
+ queryVector: Array[Double],
+ k: Int,
+ metric: DistanceMetric.Value,
+ runtimeOptions: Map[String, String] = Map.empty): LogicalPlan
+
+ /**
+ * Build a plan that finds the k nearest corpus rows for each row in the query table.
+ *
+ * @param spark active SparkSession
+ * @param corpusTable resolved corpus table
+ * @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)
+ * @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,
+ corpusTable: VectorSearchTable,
+ corpusEmbeddingCol: String,
+ queryDf: DataFrame,
+ queryEmbeddingCol: String,
+ k: Int,
+ metric: DistanceMetric.Value,
+ runtimeOptions: Map[String, String] = Map.empty): LogicalPlan
+}
+
+/**
+ * Resolves [[SearchAlgorithm]] values to [[VectorSearchAlgorithm]] implementations
+ * and provides shared validation helpers used across algorithms.
+ */
+object HoodieVectorSearchPlanBuilder {
+
+ private[analysis] val LOG = LoggerFactory.getLogger(getClass)
+
+ val DISTANCE_COL = "_hudi_distance"
+ val FILE_GROUP_ID_COL = "_hudi_file_group_id"
+ val VECTOR_ONLY_PROJECTION_OPT = "vector.exact_fetch.vector_only_projection"
+ val VECTOR_ONLY_PROJECTION_SPARK_CONF = "spark.hudi.vector.exact_fetch.vector_only_projection"
+ 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_"
+
+ private[analysis] def elapsedMs(startNs: Long): Long = (System.nanoTime() - startNs) / 1000000L
+
+ private[analysis] def summarizeInts(values: Seq[Int], limit: Int = 16): String = {
+ val preview = values.take(limit).mkString("[", ", ", "]")
+ if (values.length > limit) s"$preview...(${values.length} total)" else preview
+ }
+
+ private[analysis] def summarizeShardCounts(shardCounts: collection.Map[Int, Int], limit: Int = 16): String = {
+ val preview = shardCounts.toSeq.sortBy(_._1).take(limit).map { case (clusterId, count) =>
+ s"$clusterId->$count"
+ }.mkString("[", ", ", "]")
+ if (shardCounts.size > limit) s"$preview...(${shardCounts.size} total)" else preview
+ }
+
+ private[analysis] def logRddLineage(label: String, rdd: RDD[_]): Unit = {
+ Try(rdd.toDebugString) match {
+ case Success(debug) =>
+ LOG.info(s"[vector_search][$label] lineage:\n$debug")
+ case Failure(error) =>
+ LOG.warn(s"[vector_search][$label] unable to fetch RDD lineage", error)
+ }
+ }
+
+ /** Resolve a [[SearchAlgorithm]] enum value to its implementation. */
+ def resolveAlgorithm(algorithm: SearchAlgorithm.Value): VectorSearchAlgorithm = algorithm match {
+ case SearchAlgorithm.BRUTE_FORCE => BruteForceSearchAlgorithm
+ case SearchAlgorithm.IVF_RABITQ_MDT => IvfRaBitQMdtSearchAlgorithm
+ case other => throw new HoodieAnalysisException(
+ s"Unsupported search algorithm: $other")
+ }
+
+ 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>")
+ }
+ }
+
+ /**
+ * 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 = {
+ extractVectorSchema(df, embeddingCol).map(_.getDimension).foreach { corpusDim =>
+ if (corpusDim != queryDim) {
+ throw new HoodieAnalysisException(
+ s"Query vector dimension ($queryDim) does not match " +
+ s"corpus embedding dimension ($corpusDim) for column '$embeddingCol'")
+ }
+ }
+ }
+
+ /**
+ * 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) {
+ 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>).")
+ }
+ }
+
+ /**
+ * 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 = {
+ (extractVectorSchema(corpusDf, corpusCol).map(_.getDimension), extractVectorSchema(queryDf, queryCol).map(_.getDimension)) 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
+ }
+ }
+
+ 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 =>
+ throw new HoodieAnalysisException(
+ s"Embedding column '$colName' has type $other, expected an array type")
+ }
+ }
+
+ /** Extracts VECTOR schema from column metadata, if present. */
+ private[analysis] def extractVectorSchema(df: DataFrame, colName: String): Option[HoodieSchema.Vector] = {
+ 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)
+ 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[analysis] def quoteIdentifier(name: String): String = s"`${name.replace("`", "``")}`"
+}
+
+/**
+ * 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.
+ */
+object BruteForceSearchAlgorithm extends VectorSearchAlgorithm {
+
+ import HoodieVectorSearchPlanBuilder._
+
+ override val name: String = "brute_force"
+
+ 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 def filterOption(runtimeOptions: Map[String, String]): Option[String] =
+ runtimeOptions.get(HoodieVectorSearchTableValuedFunction.BRUTE_FORCE_FILTER_OPT)
+
+ private def maxDistanceOption(runtimeOptions: Map[String, String]): Option[Double] =
+ runtimeOptions.get(HoodieVectorSearchTableValuedFunction.BRUTE_FORCE_MAX_DISTANCE_OPT).map(_.toDouble)
+
+ override def buildSingleQueryPlan(
+ spark: SparkSession,
+ corpusTable: VectorSearchTable,
+ embeddingCol: String,
+ queryVector: Array[Double],
+ k: Int,
+ metric: DistanceMetric.Value,
+ runtimeOptions: Map[String, String] = Map.empty): LogicalPlan = {
+ val corpusDf = corpusTable.df
+ validateEmbeddingColumn(corpusDf, embeddingCol)
+ validateQueryVectorDimension(corpusDf, embeddingCol, queryVector.length)
+
+ val elemType = getElementType(corpusDf, embeddingCol)
+ var filteredDf = corpusDf.filter(col(embeddingCol).isNotNull)
Review Comment:
🤖 nit: the `var x = ...; opt.foreach(v => x = f(x, v))` pattern appears four times here (filteredDf/scored in both methods) — would `filterOption(runtimeOptions).foldLeft(corpusDf.filter(...))(applyFilter)` / `maxDistanceOption(...).foldLeft(scored)((df, d) => df.filter(col(DISTANCE_COL) <= d))` read better and let these be `val`s?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataPayload.java:
##########
@@ -253,6 +300,410 @@ protected HoodieMetadataPayload(String key, HoodieSecondaryIndexInfo secondaryIn
this(key, MetadataPartitionType.SECONDARY_INDEX.getRecordType(), null, null, null, null, secondaryIndexMetadata, secondaryIndexMetadata.getIsDeleted());
}
+ protected HoodieMetadataPayload(String key, Object vectorIndexInfo) {
+ this.key = key;
+ this.type = MetadataPartitionType.VECTOR_INDEX.getRecordType();
+ this.vectorIndexMetadata = vectorIndexInfo;
+ this.isDeletedRecord = vectorIndexInfo instanceof HoodieVectorIndexTombstone;
+ }
+
+ /**
+ * Create the singleton reader-visible generation pointer.
+ */
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexActiveManifestRecord(
+ Integer activeGeneration, String metadataPartitionPath) {
+ String recordKey = VectorIndexMetadataKey.activeManifest();
+ HoodieVectorIndexActiveManifest manifest = new HoodieVectorIndexActiveManifest(1, activeGeneration);
+ return new HoodieAvroRecord<>(
+ new HoodieKey(recordKey, metadataPartitionPath),
+ new HoodieMetadataPayload(recordKey, manifest));
+ }
+
+ /**
+ * Create the generation-one centroid record for the given index partition.
+ */
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexCentroidsRecord(
+ ByteBuffer centroidBytes, String partitionPath) {
+ HoodieVectorIndexCentroids centroids = new HoodieVectorIndexCentroids(
+ ByteBuffer.allocate(0),
+ centroidBytes,
+ ByteBuffer.allocate(0));
+ String recordKey = VectorIndexMetadataKey.centroids(1, 0);
+ HoodieMetadataPayload payload = new HoodieMetadataPayload(recordKey, centroids);
+ HoodieKey key = new HoodieKey(recordKey, partitionPath);
+ return new HoodieAvroRecord<>(key, payload);
+ }
+
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexCentroidsRecord(
+ int generation,
+ int chunk,
+ ByteBuffer clusterIds,
+ ByteBuffer centroidBytes,
+ ByteBuffer clusterRadii,
+ String partitionPath) {
+ String recordKey = VectorIndexMetadataKey.centroids(generation, chunk);
+ HoodieVectorIndexCentroids centroids = new HoodieVectorIndexCentroids(
+ clusterIds, centroidBytes, clusterRadii);
+ return new HoodieAvroRecord<>(
+ new HoodieKey(recordKey, partitionPath),
+ new HoodieMetadataPayload(recordKey, centroids));
+ }
+
+ /**
+ * Create the generation-one quantizer metadata record for the given index partition.
+ */
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexQuantizerMetadataRecord(
+ String quantizerType,
+ int quantizedCodeBytes,
+ long randomSeed,
+ boolean assumeNormalized,
+ String partitionPath) {
+ return createVectorIndexQuantizerMetadataRecord(
+ quantizerType,
+ quantizedCodeBytes,
+ 1,
+ randomSeed,
+ assumeNormalized,
+ partitionPath);
+ }
+
+ public static HoodieRecord<HoodieMetadataPayload> createVectorIndexQuantizerMetadataRecord(
+ String quantizerType,
+ int quantizedCodeBytes,
+ int rabitqBits,
+ long randomSeed,
+ boolean assumeNormalized,
+ String partitionPath) {
+ return createVectorIndexQuantizerMetadataRecord(1, 0, quantizerType, randomSeed, null, partitionPath);
Review Comment:
🤖 nit: this overload silently drops `quantizedCodeBytes`, `rabitqBits` and `assumeNormalized` — could you either remove those parameters or persist them, so callers aren't misled into thinking they're recorded?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexOptions.java:
##########
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.common.index.vector;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Options accepted by {@code CREATE INDEX ... USING VECTOR}.
+ *
+ * <p>The indexed column's Hudi {@code VECTOR(D[, elementType])} schema is authoritative for
+ * dimension and element type. Index options configure only the acceleration structure. DDL
+ * implementations must call {@link #validateAndNormalize(Map)} before persisting an index
+ * definition; individual parsing helpers are intentionally private so aggregate validation cannot
+ * be bypassed.
+ */
+public final class VectorIndexOptions {
+
+ public static final String METRIC = "vector.metric";
+ public static final String QUANTIZER = "vector.quantizer";
+ public static final String NUM_CLUSTERS = "vector.num_clusters";
+ public static final String MAX_ITER = "vector.max_iter";
+ public static final String RABITQ_BITS = "vector.rabitq.bits";
+ public static final String RABITQ_SEED = "vector.rabitq.seed";
+ public static final String RABITQ_ASSUME_NORMALIZED = "vector.rabitq.assume_normalized";
+ public static final String QUERY_NUM_PROBES = "vector.query.nprobes";
+ public static final String QUERY_REFINE_FACTOR = "vector.query.refine_factor";
+ public static final String QUERY_MODE = "vector.query.mode";
+ public static final String FRESHNESS_POLICY = "vector.freshness.policy";
+ public static final String STALE_LOCATOR_POLICY = "vector.stale.locator.policy";
+ public static final String FETCH_VERIFY_KEYS = "vector.fetch.verify.keys";
+
+ public static final VectorDistanceMetric DEFAULT_METRIC = VectorDistanceMetric.COSINE;
+ public static final VectorQuantizer DEFAULT_QUANTIZER = VectorQuantizer.IVF_RABITQ;
+ public static final int DEFAULT_NUM_CLUSTERS = 256;
+ public static final int DEFAULT_MAX_ITER = 20;
+ public static final int DEFAULT_RABITQ_BITS = 4;
+ public static final long DEFAULT_RABITQ_SEED = 42L;
+ public static final int DEFAULT_NUM_PROBES = 32;
+ public static final int DEFAULT_REFINE_FACTOR = 50;
+ public static final VectorQueryMode DEFAULT_QUERY_MODE = VectorQueryMode.EXACT_RERANK;
+ public static final VectorStalePolicy DEFAULT_EXACT_FRESHNESS_POLICY = VectorStalePolicy.FAIL;
+ public static final VectorStalePolicy DEFAULT_APPROXIMATE_FRESHNESS_POLICY = VectorStalePolicy.WARN;
+ public static final VectorStalePolicy DEFAULT_STALE_LOCATOR_POLICY = VectorStalePolicy.FALLBACK;
+ public static final boolean DEFAULT_FETCH_VERIFY_KEYS = true;
+
+ /** Parsed immutable view used by bootstrap and maintenance after aggregate validation. */
+ public static final class ResolvedOptions {
+ public final VectorDistanceMetric metric;
+ public final VectorQuantizer quantizer;
+ public final int numClusters;
+ public final int maxIterations;
+ public final int rabitqBits;
+ public final long rabitqSeed;
+ public final boolean assumeNormalized;
+
+ private ResolvedOptions(Map<String, String> options) {
+ this.metric = getMetric(options);
+ this.quantizer = getQuantizer(options);
+ this.numClusters = getNumClusters(options);
+ this.maxIterations = getMaxIter(options);
+ this.rabitqBits = getRaBitQBits(options);
+ this.rabitqSeed = getRaBitQSeed(options);
+ this.assumeNormalized = shouldAssumeNormalizedVectors(options);
+ }
+ }
+
+ /** Validates the complete map before exposing parsed writer-side settings. */
+ public static ResolvedOptions resolve(Map<String, String> options) {
+ return new ResolvedOptions(validateAndNormalize(options));
+ }
+
+ private static final Set<String> SUPPORTED_OPTIONS = Collections.unmodifiableSet(
+ new HashSet<>(Arrays.asList(
+ METRIC,
+ QUANTIZER,
+ NUM_CLUSTERS,
+ MAX_ITER,
+ RABITQ_BITS,
+ RABITQ_SEED,
+ RABITQ_ASSUME_NORMALIZED,
+ QUERY_NUM_PROBES,
+ QUERY_REFINE_FACTOR,
+ QUERY_MODE,
+ FRESHNESS_POLICY,
+ STALE_LOCATOR_POLICY,
+ FETCH_VERIFY_KEYS)));
+
+ private VectorIndexOptions() {
+ }
+
+ /**
+ * Validates the complete option map and returns canonical values for persistence.
+ *
+ * <p>The returned map contains every supported option, including explicit defaults. Unknown,
+ * retired, misspelled, and invalid options are rejected instead of being silently ignored.
+ */
+ public static Map<String, String> validateAndNormalize(Map<String, String> options) {
+ Set<String> unknownOptions = new HashSet<>(options.keySet());
+ unknownOptions.removeAll(SUPPORTED_OPTIONS);
+ if (!unknownOptions.isEmpty()) {
+ throw new IllegalArgumentException("Unsupported vector index options: " + unknownOptions);
+ }
+
+ VectorDistanceMetric metric = getMetric(options);
+ VectorQuantizer quantizer = getQuantizer(options);
+ int numClusters = getNumClusters(options);
+ int maxIter = getMaxIter(options);
+ int bits = getRaBitQBits(options);
+ long seed = getRaBitQSeed(options);
+ boolean assumeNormalized = shouldAssumeNormalizedVectors(options);
+ int numProbes = getNumProbes(options);
+ int refineFactor = getRefineFactor(options);
+ VectorQueryMode queryMode = getQueryMode(options);
+ VectorStalePolicy freshnessPolicy = getFreshnessPolicy(options);
+ VectorStalePolicy staleLocatorPolicy = getStaleLocatorPolicy(options);
+ boolean verifyFetchKeys = shouldVerifyFetchKeys(options);
+
+ if (numProbes > numClusters) {
+ throw new IllegalArgumentException(
+ "Option '" + QUERY_NUM_PROBES + "' must not exceed '" + NUM_CLUSTERS + "': "
+ + numProbes + " > " + numClusters);
+ }
+
+ Map<String, String> normalized = new LinkedHashMap<>();
+ normalized.put(METRIC, metric.name().toLowerCase(Locale.ROOT));
+ normalized.put(QUANTIZER, quantizer.name());
+ normalized.put(NUM_CLUSTERS, String.valueOf(numClusters));
+ normalized.put(MAX_ITER, String.valueOf(maxIter));
+ normalized.put(RABITQ_BITS, String.valueOf(bits));
+ normalized.put(RABITQ_SEED, String.valueOf(seed));
+ normalized.put(RABITQ_ASSUME_NORMALIZED, String.valueOf(assumeNormalized));
+ normalized.put(QUERY_NUM_PROBES, String.valueOf(numProbes));
+ normalized.put(QUERY_REFINE_FACTOR, String.valueOf(refineFactor));
+ normalized.put(QUERY_MODE, queryMode.name().toLowerCase(Locale.ROOT));
+ normalized.put(FRESHNESS_POLICY, freshnessPolicy.name().toLowerCase(Locale.ROOT));
+ normalized.put(STALE_LOCATOR_POLICY, staleLocatorPolicy.name().toLowerCase(Locale.ROOT));
+ normalized.put(FETCH_VERIFY_KEYS, String.valueOf(verifyFetchKeys));
+ return Collections.unmodifiableMap(normalized);
+ }
+
+ public static VectorDistanceMetric getMetric(Map<String, String> options) {
+ String value = getOption(options, METRIC, DEFAULT_METRIC.name());
+ try {
+ return VectorDistanceMetric.fromString(value);
+ } catch (IllegalArgumentException e) {
+ throw unsupportedValue(METRIC, value, e);
+ }
+ }
+
+ private static VectorQuantizer getQuantizer(Map<String, String> options) {
+ String value = getOption(options, QUANTIZER, DEFAULT_QUANTIZER.name());
+ try {
+ return VectorQuantizer.fromString(value);
+ } catch (IllegalArgumentException e) {
+ throw unsupportedValue(QUANTIZER, value, e);
+ }
+ }
+
+ private static int getNumClusters(Map<String, String> options) {
+ return getPositiveInt(options, NUM_CLUSTERS, DEFAULT_NUM_CLUSTERS);
+ }
+
+ private static int getMaxIter(Map<String, String> options) {
+ return getPositiveInt(options, MAX_ITER, DEFAULT_MAX_ITER);
+ }
+
+ private static int getRaBitQBits(Map<String, String> options) {
+ int bits = getInt(options, RABITQ_BITS, DEFAULT_RABITQ_BITS);
+ if (bits < 1 || bits > 8) {
+ throw new IllegalArgumentException(
+ "Option '" + RABITQ_BITS + "' must be between 1 and 8: " + bits);
+ }
+ return bits;
+ }
+
+ private static long getRaBitQSeed(Map<String, String> options) {
+ String value = getOption(options, RABITQ_SEED, String.valueOf(DEFAULT_RABITQ_SEED));
+ try {
+ return Long.parseLong(value);
+ } catch (NumberFormatException e) {
+ throw invalidNumber(RABITQ_SEED, value, e);
+ }
+ }
+
+ private static boolean shouldAssumeNormalizedVectors(Map<String, String> options) {
+ return getBoolean(options, RABITQ_ASSUME_NORMALIZED, false);
+ }
+
+ public static int getNumProbes(Map<String, String> options) {
+ return getPositiveInt(options, QUERY_NUM_PROBES, DEFAULT_NUM_PROBES);
+ }
+
+ public static int getRefineFactor(Map<String, String> options) {
+ return getPositiveInt(options, QUERY_REFINE_FACTOR, DEFAULT_REFINE_FACTOR);
+ }
+
+ private static VectorQueryMode getQueryMode(Map<String, String> options) {
+ String value = getOption(options, QUERY_MODE, DEFAULT_QUERY_MODE.name());
+ try {
+ return VectorQueryMode.fromString(value);
+ } catch (IllegalArgumentException e) {
+ throw unsupportedValue(QUERY_MODE, value, e);
+ }
+ }
+
+ public static VectorStalePolicy getFreshnessPolicy(Map<String, String> options) {
+ VectorStalePolicy defaultPolicy = getQueryMode(options) == VectorQueryMode.APPROXIMATE
+ ? DEFAULT_APPROXIMATE_FRESHNESS_POLICY : DEFAULT_EXACT_FRESHNESS_POLICY;
+ return getStalePolicy(options, FRESHNESS_POLICY, defaultPolicy);
+ }
+
+ public static VectorStalePolicy getStaleLocatorPolicy(Map<String, String> options) {
+ return getStalePolicy(options, STALE_LOCATOR_POLICY, DEFAULT_STALE_LOCATOR_POLICY);
+ }
+
+ private static VectorStalePolicy getStalePolicy(
+ Map<String, String> options, String key, VectorStalePolicy defaultPolicy) {
+ String value = getOption(options, key, defaultPolicy.name());
+ try {
+ return VectorStalePolicy.fromString(value);
+ } catch (IllegalArgumentException e) {
+ throw unsupportedValue(key, value, e);
+ }
+ }
+
+ public static boolean isApproximateSearchMode(Map<String, String> options) {
+ return getQueryMode(options) == VectorQueryMode.APPROXIMATE;
+ }
+
+ /** The current posting format uses symmetric scoring. */
+ public static boolean isRaBitQAsymmetricScoring(Map<String, String> options) {
+ return false;
Review Comment:
🤖 nit: `isRaBitQAsymmetricScoring` and `isFinalistArbiterEnabled` take `options` but ignore it and return constants — could these be plain constants (or drop the parameter) so readers don't go looking for a config key that doesn't exist?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-common/src/main/java/org/apache/hudi/metadata/HoodieTableMetadataUtil.java:
##########
@@ -213,6 +215,71 @@ public class HoodieTableMetadataUtil {
HoodieRecord.HoodieMetadataField.PARTITION_PATH_METADATA_FIELD.getFieldName(),
HoodieRecord.HoodieMetadataField.COMMIT_TIME_METADATA_FIELD.getFieldName()));
+ public static String getVectorIndexGenerationManifestKey(int generationId) {
+ return VectorIndexMetadataKey.manifest(generationId);
+ }
+
+ public static boolean isVectorIndexGenerationManifestKey(String recordKey) {
+ return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_MANIFEST);
+ }
+
+ public static boolean isVectorIndexQuantizerKey(String recordKey) {
+ return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_QUANTIZER);
+ }
+
+ public static boolean isVectorIndexCentroidsKey(String recordKey) {
+ return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_CENTROIDS);
+ }
+
+ public static String getVectorIndexClusterKey(int generationId, int clusterId) {
+ return VectorIndexMetadataKey.clusterStats(generationId, clusterId);
+ }
+
+ public static boolean isVectorIndexClusterKey(String recordKey) {
+ return hasVectorIndexFamily(recordKey, VectorIndexMetadataKey.FAMILY_CLUSTER_STATS);
+ }
+
+ public static String getVectorIndexPostingKey(int generationId, int clusterId, int shardId, String recordKey) {
+ return VectorIndexMetadataKey.postingDelta(generationId, clusterId, shardId, recordKey);
+ }
+
+ public static String getVectorIndexPostingPrefix(int generationId, int clusterId) {
+ return VectorIndexMetadataKey.postingPrefix(generationId, clusterId, 0).substring(0, 9);
Review Comment:
🤖 nit: the `substring(0, 9)` cluster-prefix trick is repeated here, in `getVectorIndexPostingRoutingKey`, and in `VectorPostingPrefixRawKey` — would it make sense to add a `VectorIndexMetadataKey.postingClusterPrefix(generation, clusterId)` so the 1+4+4 byte layout lives in one place?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorExactSearchPlanner.scala:
##########
@@ -0,0 +1,390 @@
+/*
+ * 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.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.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, 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}
+
+/** Exact reranking and file-slice resolution for the IVF/RaBitQ search plan. */
+private[analysis] object VectorExactSearchPlanner extends SparkAdapterSupport {
+ import HoodieVectorSearchPlanBuilder._
+ import IvfRaBitQMdtSearchAlgorithm.{isLogResidentCandidate, isPartitionedRecordIndex, toVectorMetric, VectorSearchPlanContext}
+
+ private val LOG = LoggerFactory.getLogger(getClass)
+
+ def buildRaBitQExactCandidateDf(
+ spark: SparkSession,
+ ctx: VectorSearchPlanContext,
+ outputSchema: StructType,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Float],
+ exactQueryVector: Array[Double],
+ k: Int): DataFrame = {
+ val planStartNs = System.nanoTime()
+ val storageConf = HadoopFSUtils.getStorageConf(spark.sessionState.newHadoopConf())
+ val cacheLoadStartNs = 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 topClusters = cache.findTopClusters(queryVector, numProbes, metric)
+ val topClusterSeq = topClusters.toSeq
+ val cacheLoadMs = elapsedMs(cacheLoadStartNs)
+
+ // Shard counts from cache — no extra MDT IO.
+ val shardCounts = cache.getShardCounts(topClusters)
+ val shardCountsScala = shardCounts.asScala.map { case (clusterId, count) =>
+ clusterId.intValue() -> count.intValue()
+ }.toMap
+ val postingPrefixCount = VectorIndexMdtSearchUtils.buildPostingPrefixes(cache.getGenerationId, shardCounts).size()
+ LOG.info(
+ s"[vector_search][plan] basePath=${ctx.basePath} indexPartition=${ctx.indexPartition} cacheLoadMs=$cacheLoadMs " +
+ s"generationId=${cache.getGenerationId} dimension=${cache.getDimension} numClusters=${cache.numClusters()} " +
+ s"numProbes=$numProbes refineFactor=${VectorIndexOptions.getRefineFactor(mergedOptions.asJava)} residualEncoding=$residualEncoding " +
+ s"topClusters=${summarizeInts(topClusterSeq)} shardCounts=${summarizeShardCounts(shardCountsScala)} " +
+ s"postingPrefixes=$postingPrefixCount")
+ if (shardCounts.isEmpty) {
+ sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+ } else {
+ val refineK = math.max(k, k * VectorIndexOptions.getRefineFactor(mergedOptions.asJava))
+ val emptyDf = sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+
+ 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, refineK)
+ val topCandidateMaterializeStartNs = System.nanoTime()
+ val materializedCandidates = try {
+ topCandidates.collectAsList().asScala.toSeq
+ } finally {
+ topCandidates.unpersistWithDependencies()
+ }
+ val topCandidateList = materializedCandidates
+ .filter(candidate => candidate.getFileGroupId != null && candidate.getFileGroupId.nonEmpty)
+ val candidatesWithoutLocator = materializedCandidates.size - topCandidateList.size
+ LOG.info(
+ s"[vector_search][stage][materialize_top_candidates] refineK=$refineK rerankCandidates=${materializedCandidates.size} " +
+ s"candidatesWithLocators=${topCandidateList.size} candidatesWithoutLocators=$candidatesWithoutLocator " +
+ s"distinctFileGroups=${topCandidateList.map(_.getFileGroupId).distinct.size} elapsedMs=${elapsedMs(topCandidateMaterializeStartNs)}")
+ if (candidatesWithoutLocator > 0) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search found $candidatesWithoutLocator rerank candidates without file-group locators. " +
+ s"Exact rerank requires per-candidate posting locators.")
+ }
+
+ // RFC-109 RLI finalist arbiter. Exact mode uses refreshed RLI locations for every live
+ // candidate: SERVE retains positional trust; STALE is resolved by key in the live base file.
+ val arbiterEnabled = VectorIndexOptions.isFinalistArbiterEnabled(mergedOptions.asJava)
+ val serveCandidates: Seq[ScoredVectorPostingMatch] =
+ if (arbiterEnabled) {
+ val arb = VectorIndexMdtSearchUtils.arbitrateMaterializedFinalists(
+ ctx.metadataTable, topCandidateList.asJava, isPartitionedRecordIndex(ctx))
+ val staleFailPolicy = VectorIndexOptions.isStaleLocatorPolicyFail(mergedOptions.asJava)
+ LOG.info(
+ s"[vector_search][exact][arbiter] arbiterExclusions{stale=${arb.staleCount}, deleted=${arb.deletedCount}} " +
+ s"serve=${arb.serve.size} stalePolicy=${if (staleFailPolicy) "fail" else "fallback"}")
+ if (staleFailPolicy && !arb.stale.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search rejected ${arb.staleCount} stale finalist(s) because " +
+ "vector.stale.locator.policy=FAIL.")
+ }
+ arb.serve.asScala.toSeq ++ arb.stale.asScala.toSeq
+ } else {
+ topCandidateList
+ }
+
+ if (serveCandidates.isEmpty) {
+ emptyDf
+ } else {
+ val candidatesByPartition = serveCandidates
+ .groupBy(candidate => Option(candidate.getLocation).map(_.getPartitionPath).getOrElse(""))
+ .map { case (partitionPath, candidates) =>
+ partitionPath -> candidates.flatMap(candidate => Option(candidate.getLocation).map(_.getFileId)).toSet
+ }
+ val fileSliceMap = preResolveFileSlices(
+ ctx.metadataTable, ctx.metaClient, ctx.timeline, ctx.latestInstantTime, candidatesByPartition)
+ if (fileSliceMap.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search resolved zero file slices for ${serveCandidates.size} rerank candidates " +
+ s"across ${candidatesByPartition.values.map(_.size).sum} file groups. Candidate locators are present, so this indicates " +
+ s"a file-slice resolution bug or stale vector index.")
+ }
+ val bFileSliceMap = spark.sparkContext.broadcast(fileSliceMap)
+ // MOR log-resident candidates carry their raw embedding in a log block, not the base Parquet
+ // file, so the positional fetcher cannot read them. Split them out and materialize them via a
+ // merged (base + log) read; base-resident candidates keep the fast positional path.
+ val (logResidentCandidates, baseCandidates) = serveCandidates.partition { candidate =>
+ bFileSliceMap.value.get(candidate.getLocation.getFileId).exists { slice =>
+ isLogResidentCandidate(candidate, slice.getBaseInstantTime, slice.getLogFiles.findAny().isPresent)
+ }
+ }
+ val exactFetchParallelism = math.max(
+ 1,
+ math.min(refineK, math.max(1, spark.sparkContext.defaultParallelism * 4)))
+ val groupedCandidates = baseCandidates.groupBy(_.getLocation.getFileId).toSeq.sortBy(_._1)
+ val targetBatchesPerGroup = math.max(1, exactFetchParallelism / math.max(1, groupedCandidates.size))
+ val fetchBatches = groupedCandidates.flatMap { case (fgId, candidates) =>
+ val sorted = candidates.sortBy(candidate => candidate.getLocation.getPosition)
+ val batchSize = math.max(1, math.ceil(sorted.size.toDouble / targetBatchesPerGroup).toInt)
+ sorted.grouped(batchSize).map(batch => (fgId, batch)).toSeq
+ }
+ val byFetchBatch: RDD[(String, Seq[ScoredVectorPostingMatch])] =
+ if (fetchBatches.isEmpty) {
+ spark.sparkContext.emptyRDD[(String, Seq[ScoredVectorPostingMatch])]
+ } else {
+ spark.sparkContext.parallelize(fetchBatches, math.min(exactFetchParallelism, fetchBatches.size))
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read_plan] rerankCandidates=${serveCandidates.size} " +
+ s"candidateGroups=${groupedCandidates.size} resolvedSlices=${fileSliceMap.size} " +
+ s"fileGroupPartitions=${byFetchBatch.getNumPartitions} exactFetchParallelism=$exactFetchParallelism")
+
+ val vectorOnlyProjection = mergedOptions.get(VECTOR_ONLY_PROJECTION_OPT)
+ .orElse(spark.conf.getOption(VECTOR_ONLY_PROJECTION_SPARK_CONF))
+ .exists(_.equalsIgnoreCase("true"))
+ val fetchSkippedLogResidentAcc: LongAccumulator =
+ spark.sparkContext.longAccumulator("vector_fetch_skipped_log_resident")
+ val fetchKeyMismatchAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_fetch_key_mismatch")
+ val rowsMaterializedAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_rows_materialized")
+ val exactRows: RDD[InternalRow] = byFetchBatch.mapPartitions { batches =>
+ val startNs = System.nanoTime()
+ val rows = ArrayBuffer.empty[InternalRow]
+ var batchCount = 0
+ var resolvedSlices = 0
+ var missingSlices = 0
+ var candidateKeys = 0L
+ var baseFiles = 0
+ var logFiles = 0L
+ var totalSliceBytes = 0L
+ var staleCandidates = 0L
+ var fetchSkippedLogResident = 0L
+ var fetchKeyMismatch = 0L
+ var keyLookupRowsDecoded = 0L
+ var metrics = PositionalParquetFetcher.FetchMetrics()
+ val fetcher = new PositionalParquetFetcher(
+ storageConf.newInstance().unwrap(),
+ embeddingCol,
+ vectorSchema,
+ exactQueryVector,
+ ctx.metric,
+ outputSchema,
+ mergedOptions.getOrElse("vector.exact_fetch.offset_index_missing_policy", "fallback").equalsIgnoreCase("fail"),
+ vectorOnlyProjection,
+ VectorIndexOptions.shouldVerifyFetchKeys(mergedOptions.asJava))
+ val keyLocator = new ParquetRecordKeyLocator(storageConf.newInstance().unwrap())
+ while (batches.hasNext) {
+ val (fgId, candidates) = batches.next()
+ batchCount += 1
+ val fileSliceOpt = bFileSliceMap.value.get(fgId)
+ if (fileSliceOpt.isEmpty) {
+ missingSlices += 1
+ throw new HoodieAnalysisException(
+ s"Vector exact search could not resolve file group '$fgId' for ${candidates.size} rerank candidates. " +
+ s"No record-level-index fallback is active in this path.")
+ } else {
+ resolvedSlices += 1
+ val fileSlice = fileSliceOpt.get
+ val logFileCount = fileSlice.getLogFiles.count()
+ val currentBaseInstantTime = fileSlice.getBaseInstantTime
+ val logResident = candidates.filter(candidate =>
+ isLogResidentCandidate(candidate, currentBaseInstantTime, logFileCount > 0))
+ if (logResident.nonEmpty) {
+ fetchSkippedLogResident += logResident.size
+ fetchSkippedLogResidentAcc.add(logResident.size)
+ LOG.warn(
+ s"Vector exact fetch skipped ${logResident.size} log-resident candidate(s) in file group '$fgId'; " +
+ "they remain a result shortfall and the retained candidate pool will refill where candidates remain")
+ }
+ val liveCandidates = candidates.filterNot(logResident.toSet)
+ staleCandidates += liveCandidates.count(candidate =>
+ candidate.getArbiterDecision == VectorIndexArbiter.Decision.STALE)
+ if (liveCandidates.nonEmpty) {
+ val positionalCandidates = liveCandidates.filter { candidate =>
+ val location = candidate.getLocation
+ location != null && location.getPosition >= 0 &&
+ candidate.getArbiterDecision != VectorIndexArbiter.Decision.STALE &&
+ (location.getInstantTime == null || location.getInstantTime == currentBaseInstantTime)
+ }
+ val keyCandidates = liveCandidates.filterNot(positionalCandidates.toSet)
+ val recordKeys = liveCandidates.map(_.getRecordKey).toSet
+ candidateKeys += recordKeys.size
+ if (fileSlice.getBaseFile.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact key fallback requires a base Parquet file for file group '$fgId'; " +
+ "the remaining candidates are not classified as log-resident.")
+ }
+ baseFiles += 1
+ logFiles += logFileCount
+ totalSliceBytes += fileSlice.getTotalFileSize
+ val baseFilePath = fileSlice.getBaseFile.get.getPath
+ val locatedByKey = keyLocator.locate(baseFilePath, keyCandidates)
+ keyLookupRowsDecoded += locatedByKey.metrics.rowsDecoded
+ if (locatedByKey.metrics.located != keyCandidates.size) {
+ throw new HoodieAnalysisException(
+ s"Vector exact key fallback located ${locatedByKey.metrics.located} of ${keyCandidates.size} " +
+ s"candidate key(s) in the live base file for file group '$fgId'.")
+ }
+ val firstFetch = fetcher.fetch(baseFilePath, positionalCandidates ++ locatedByKey.candidates)
+ rows ++= firstFetch.rows
+ metrics = metrics.add(firstFetch.metrics)
+ rowsMaterializedAcc.add(firstFetch.metrics.rowsMaterialized)
+ if (firstFetch.retryByKey.nonEmpty) {
+ fetchKeyMismatch += firstFetch.retryByKey.size
+ fetchKeyMismatchAcc.add(firstFetch.retryByKey.size)
+ val retries = keyLocator.locate(baseFilePath, firstFetch.retryByKey)
+ keyLookupRowsDecoded += retries.metrics.rowsDecoded
+ if (retries.metrics.located != firstFetch.retryByKey.size) {
+ throw new HoodieAnalysisException(
+ s"Vector exact fetch could not relocate ${firstFetch.retryByKey.size - retries.metrics.located} " +
+ s"key mismatch(es) in file group '$fgId'.")
+ }
+ val retryFetch = fetcher.fetch(baseFilePath, retries.candidates)
+ if (retryFetch.retryByKey.nonEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact fetch could not self-heal ${retryFetch.retryByKey.size} key mismatch(es) " +
+ s"in file group '$fgId' after key-based relocation.")
+ }
+ rows ++= retryFetch.rows
+ metrics = metrics.add(retryFetch.metrics)
+ rowsMaterializedAcc.add(retryFetch.metrics.rowsMaterialized)
+ }
+ }
+ }
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read] rerankCandidates=$candidateKeys candidateFileCount=$resolvedSlices " +
+ s"fetchBatches=$batchCount resolvedSlices=$resolvedSlices missingSlices=$missingSlices " +
+ s"rowsDecoded=${metrics.rowsDecoded} rowsMaterialized=${metrics.rowsMaterialized} " +
+ s"keyLookupRowsDecoded=$keyLookupRowsDecoded fetchKeyMismatch=$fetchKeyMismatch " +
+ s"fetchSkippedLogResident=$fetchSkippedLogResident " +
+ s"rowGroupsTotal=${metrics.rowGroupsTotal} rowGroupsSelected=${metrics.rowGroupsSelected} " +
+ s"pagesInSelectedRowGroups=${metrics.pagesInSelectedRowGroups} pagesSelected=${metrics.pagesSelected} " +
+ s"offsetIndexHits=${metrics.offsetIndexHits} offsetIndexMissing=${metrics.offsetIndexMissing} " +
+ s"rangedGets=${metrics.rangedGets} pageBytesFetched=${metrics.pageBytesFetched} " +
+ s"physicalBytesRead=${metrics.physicalBytesRead} footerReads=${metrics.footerReads} footerCacheHits=${metrics.footerCacheHits} " +
+ s"fetchKeyMismatchMetric=${metrics.fetchKeyMismatch} staleCandidates=$staleCandidates " +
+ s"baseFiles=$baseFiles logFiles=$logFiles " +
+ s"sliceBytes=$totalSliceBytes footerMs=${metrics.footerMs} offsetIndexMs=${metrics.offsetIndexMs} " +
+ s"fetchWaitMs=${metrics.fetchWaitMs} decodeMs=${metrics.decodeMs} scoreMs=${metrics.scoreMs} elapsedMs=${elapsedMs(startNs)}")
+ rows.iterator
+ }
+ val logResidentRows: Seq[InternalRow] =
+ if (logResidentCandidates.isEmpty) {
+ Seq.empty
+ } else {
+ val dataSchema = new TableSchemaResolver(ctx.metaClient).getTableSchema(true)
+ val scorer = new VectorExactScorer(vectorSchema, exactQueryVector, ctx.metric)
+ val logFetcher = new LogResidentVectorFetcher(
+ storageConf, ctx.metaClient, dataSchema, ctx.latestInstantTime,
+ embeddingCol, outputSchema, scorer)
+ val fetched = logResidentCandidates
+ .groupBy(candidate =>
+ (Option(candidate.getLocation.getPartitionPath).getOrElse(""), candidate.getLocation.getFileId))
+ .flatMap { case ((partitionPath, fgId), candidates) =>
+ logFetcher.fetchSlice(partitionPath, fileSliceMap(fgId), candidates.map(_.getRecordKey).toSet)
Review Comment:
🤖 non-blocking: this does a full merged read of each file slice on the driver, so the cost is bounded by base+log size per touched file group rather than by refineK as the fetcher's doc says; since the slice map is already broadcast and the scorer is serializable, would it make sense to run this inside the executor `mapPartitions` path next to the positional fetch?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metadata/SparkVectorIndexBootstrap.java:
##########
@@ -0,0 +1,531 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.metadata;
+
+import org.apache.hudi.common.data.HoodieData;
+import org.apache.hudi.common.index.vector.PostingBlockBuilder;
+import org.apache.hudi.common.index.vector.RaBitQEncoder;
+import org.apache.hudi.common.index.vector.VectorDistanceMetric;
+import org.apache.hudi.common.index.vector.VectorIndexBootstrapUtils;
+import org.apache.hudi.common.index.vector.VectorIndexOptions;
+import org.apache.hudi.common.model.HoodieIndexDefinition;
+import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.schema.HoodieSchema;
+import org.apache.hudi.data.HoodieJavaRDD;
+import org.apache.hudi.spark.index.vector.TwoLevelKMeansBootstrap$;
+
+import lombok.extern.slf4j.Slf4j;
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.broadcast.Broadcast;
+import org.apache.spark.sql.SparkSession;
+
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import scala.Tuple2;
+
+/**
+ * Two-phase vector index bootstrap for Spark.
+ *
+ * <p>Phase 1: Sample vectors → train centroids with Spark ML KMeans. Training
+ * still uses a bounded sample (FAISS, ScaNN, and Milvus all use sample-based
+ * centroid training for IVF), but the iterative KMeans work stays distributed
+ * instead of collecting the sample to the driver.
+ *
+ * <p>Phase 2: Broadcast centroids → single-pass {@code mapPartitions} that
+ * assigns each vector to its nearest centroid, computes RaBitQ binary codes,
+ * and emits posting metadata records with per-candidate locators. Cluster
+ * statistics (populations, shard counts) are gathered via a lightweight
+ * {@code countByKey} before the main encoding pass.
+ *
+ * <p>Compared to the previous driver-local KMeans approach:
+ * <ul>
+ * <li>No reflection — direct {@link HoodieMetadataPayload} method calls</li>
+ * <li>Proper {@code unpersist()} in try/finally</li>
+ * <li>generation ordinal allocated from existing manifest keys</li>
+ * </ul>
+ */
+@Slf4j
+public final class SparkVectorIndexBootstrap {
+
+ /** Hard cap for percentile-driven training samples. */
+ private static final int MAX_PERCENTILE_TRAINING_SAMPLE = 10_000_000;
+
+ /** Absolute minimum sample floor for production index quality. */
+ private static final int MIN_TRAINING_SAMPLE = 1_000_000;
+
+ /** Minimum sample multiplier per cluster. FAISS recommends 256×K minimum. */
+ private static final int MIN_SAMPLE_PER_CLUSTER = 256;
+
+ /** Lower bound for percentile-driven sampling: 0.5% of total vectors. */
+ private static final double MIN_SAMPLE_PERCENT = 0.005d;
+
+ /** Upper bound for percentile-driven sampling: 1% of total vectors. */
+ private static final double MAX_SAMPLE_PERCENT = 0.01d;
+
+ private SparkVectorIndexBootstrap() {
+ }
+
+ /**
+ * Run the two-phase vector index bootstrap and return all MDT records.
+ *
+ * @param jsc Spark context
+ * @param vectorRows RDD of (recordKey, partitionPath, fileId, vectorBytes) — already
+ * read from base-table file slices
+ * @param indexDef the index definition
+ * @param vectorType resolved vector element type from the table schema
+ * @param generation vector generation ordinal
+ * @param lastUpdatedTs timestamp for MDT bookkeeping
+ * @return all metadata records to write into the vector index MDT partition
+ */
+ public static HoodieData<HoodieRecord> bootstrap(JavaSparkContext jsc,
+ JavaRDD<VectorRow> vectorRows,
+ HoodieIndexDefinition indexDef,
+ HoodieSchema.Vector.VectorElementType vectorType,
+ int dimension,
+ int generation,
+ long lastUpdatedTs) {
+ Map<String, String> options = indexDef.getIndexOptions();
+ String indexName = indexDef.getIndexName();
+ String vectorColumn = indexDef.getSourceFields().isEmpty() ? "" : indexDef.getSourceFields().get(0);
+ VectorIndexOptions.ResolvedOptions resolvedOptions = VectorIndexOptions.resolve(options);
+ int numClusters = resolvedOptions.numClusters;
+ VectorDistanceMetric metric = resolvedOptions.metric;
+
+ String quantizerType = resolvedOptions.quantizer.name();
+ long quantizerSeed = resolvedOptions.rabitqSeed;
+ int rabitqBits = resolvedOptions.rabitqBits;
+ boolean assumeNormalized = resolvedOptions.assumeNormalized;
+ boolean storeInMdt = true;
+ int targetRowsPerShard = 100_000;
Review Comment:
🤖 nit: `storeInMdt` is hardcoded to `true`, so the `else` branches below are dead — could you drop the flag (and the `residualEncoding`/`targetRowsPerShard`/`maxShardsPerCluster` locals into named constants) until there's a real alternative path?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/PositionalParquetFetcher.scala:
##########
@@ -0,0 +1,587 @@
+/*
+ * 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.spark.sql.hudi.analysis
+
+import org.apache.hudi.common.index.vector.{ScoredVectorPostingMatch, VectorIndexMdtSearchUtils, VectorPostingMatch}
+import org.apache.hudi.common.model.HoodieRecord
+import org.apache.hudi.common.schema.HoodieSchema
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileSystem, Path}
+import org.apache.parquet.column.page.PageReadStore
+import org.apache.parquet.hadoop.ParquetFileReader
+import org.apache.parquet.hadoop.metadata.{BlockMetaData, ColumnChunkMetaData, ParquetMetadata}
+import org.apache.parquet.internal.column.columnindex.OffsetIndex
+import org.apache.parquet.internal.filter2.columnindex.RowRanges
+import org.apache.parquet.io.{ColumnIOFactory, RecordReader}
+import org.apache.parquet.io.api.{Binary, Converter, GroupConverter, PrimitiveConverter, RecordMaterializer}
+import org.apache.parquet.schema.MessageType
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction.DistanceMetric
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+import org.slf4j.LoggerFactory
+
+import java.util.PrimitiveIterator
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+/**
+ * Fetches exact vector-search rerank candidates from Parquet by physical row position.
+ *
+ * The fetcher uses Parquet page indexes to turn candidate row positions into row ranges, so
+ * the reader can skip non-candidate pages instead of scanning whole Hudi file slices.
+ */
+private[analysis] final class PositionalParquetFetcher(
+ conf: Configuration,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Double],
+ metric: DistanceMetric.Value,
+ outputSchema: StructType,
+ failOnMissingOffsetIndex: Boolean,
+ vectorOnlyProjection: Boolean,
+ verifyKeys: Boolean) extends Serializable {
+
+ import PositionalParquetFetcher._
+
+ private val distanceOrdinal = outputSchema.fieldIndex(HoodieVectorSearchPlanBuilder.DISTANCE_COL)
+ private val recordKeyOrdinal = fieldOrdinal(HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ private val partitionPathOrdinal = fieldOrdinal(HoodieRecord.PARTITION_PATH_METADATA_FIELD)
+ private val fileGroupIdOrdinal = fieldOrdinal(HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL)
+ private val scorer = new VectorExactScorer(vectorSchema, queryVector, metric)
+ private var scoreMs: Long = 0L
+
+ def fetch(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch]): FetchResult = {
+ val startNs = System.nanoTime()
+ if (candidates.isEmpty) {
+ FetchResult(Iterator.empty, Seq.empty, FetchMetrics(elapsedMs = elapsedMs(startNs)))
+ } else {
+ fetchNonEmpty(filePath, candidates, startNs)
+ }
+ }
+
+ private def fetchNonEmpty(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch],
+ startNs: Long): FetchResult = {
+ val path = new Path(filePath)
+ configureRandomRead(conf)
+ val fs = path.getFileSystem(conf)
+ val fileLen = fs.getFileStatus(path).getLen
+ val fileCacheKey = s"$filePath#$fileLen"
+ val stats = Option(FileSystem.getStatistics(path.toUri.getScheme, fs.getClass))
+ val physicalBefore = stats.map(_.getBytesRead).getOrElse(0L)
+
+ val footerStartNs = System.nanoTime()
+ val footerResult = footer(conf, path, fileCacheKey)
+ val footerMs = elapsedMs(footerStartNs)
+ val metadata = footerResult.metadata
+ val fileSchema = metadata.getFileMetaData.getSchema
+ val requestedParquetSchema = requestedSchema(fileSchema)
+
+ val rowGroups = metadata.getBlocks.asScala.toIndexedSeq
+ val rowGroupStarts = cumulativeRowGroupStarts(rowGroups)
+ val byRowGroup: Map[Int, Seq[ScoredVectorPostingMatch]] =
+ candidates.groupBy(candidate => rowGroupOrdinal(rowGroupStarts, rowGroups, rowPosition(candidate)))
+
+ val rows = mutable.ArrayBuffer.empty[InternalRow]
+ val keyMismatches = mutable.ArrayBuffer.empty[ScoredVectorPostingMatch]
+ var rowGroupsSelected = 0
+ var pagesInSelectedRowGroups = 0L
+ var pagesSelected = 0L
+ var offsetIndexHits = 0L
+ var offsetIndexMissing = 0L
+ var pageBytesFetched = 0L
+ var rangedGets = 0L
+ var rowsDecoded = 0L
+ var rowsMaterialized = 0L
+ var offsetIndexMs = 0L
+ var decodeMs = 0L
+ scoreMs = 0L
+
+ val reader = new ParquetFileReader(conf, path, metadata)
+ try {
+ reader.setRequestedSchema(requestedParquetSchema)
+ byRowGroup.toSeq.sortBy(_._1).foreach { case (rgOrdinal, rgCandidates) =>
+ val block = rowGroups(rgOrdinal)
+ val rowGroupStart = rowGroupStarts(rgOrdinal)
+ val embeddingChunk = findColumn(block, embeddingCol)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet file '$filePath' is missing vector column '$embeddingCol'."))
+
+ val offsetStartNs = System.nanoTime()
+ val embeddingOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, embeddingChunk)
+ offsetIndexMs += elapsedMs(offsetStartNs)
+ val relativeRows = rgCandidates.map(candidate => rowPosition(candidate) - rowGroupStart).sorted
+ val candidateRows = relativeRows.toSet
+ val candidateByRelativeRow = rgCandidates
+ .map(candidate => rowPosition(candidate) - rowGroupStart -> candidate)
+ .toMap
+
+ if (embeddingOffsetIndex == null && failOnMissingOffsetIndex) {
+ offsetIndexMissing += 1
+ throw new IllegalStateException(
+ s"Parquet file '$filePath' has no OffsetIndex for vector column '$embeddingCol'; positional exact fetch cannot honor pages.")
+ }
+
+ rowGroupsSelected += 1
+ val (pageStore, rowRanges) = if (embeddingOffsetIndex == null) {
+ offsetIndexMissing += 1
+ LOG.warn(
+ "Parquet file {} rowGroup={} has no OffsetIndex for vector column {}; falling back to row-group fetch.",
+ filePath,
+ Int.box(rgOrdinal),
+ embeddingCol)
+ pagesInSelectedRowGroups += 0L
+ pagesSelected += 0L
+ rangedGets += 1L
+ pageBytesFetched += requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .map(_.getTotalSize)
+ .sum
+ (reader.readRowGroup(rgOrdinal), RowRanges.createSingle(block.getRowCount))
+ } else {
+ offsetIndexHits += 1
+ pagesInSelectedRowGroups += embeddingOffsetIndex.getPageCount
+ val selectedVectorPages = selectedPages(embeddingOffsetIndex, block.getRowCount, relativeRows)
+ pagesSelected += selectedVectorPages.size
+ val selectedRanges = requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .flatMap { chunk =>
+ val oiStartNs = System.nanoTime()
+ val chunkOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, chunk)
+ offsetIndexMs += elapsedMs(oiStartNs)
+ if (chunkOffsetIndex == null) {
+ offsetIndexMissing += 1
+ Seq.empty
+ } else {
+ offsetIndexHits += 1
+ selectedPages(chunkOffsetIndex, block.getRowCount, relativeRows).map { pageOrdinal =>
+ val start = chunkOffsetIndex.getOffset(pageOrdinal)
+ PageRange(start, start + chunkOffsetIndex.getCompressedPageSize(pageOrdinal))
+ }
+ }
+ }
+ val mergedRanges = mergeRanges(selectedRanges.sortBy(_.start).toSeq)
+ rangedGets += mergedRanges.size
+ pageBytesFetched += mergedRanges.iterator.map(range => range.end - range.start).sum
+ val ranges = RowRanges.create(block.getRowCount, primitiveIntIterator(selectedVectorPages), embeddingOffsetIndex)
+ (reader.readFilteredRowGroup(rgOrdinal, ranges), ranges)
+ }
+
+ val decodeStartNs = System.nanoTime()
+ val recordReader = parquetRecordReader(fileSchema, requestedParquetSchema, pageStore)
+ val rowIterator = rowRanges.iterator()
+ while (rowIterator.hasNext) {
+ val relativeRow = rowIterator.nextLong()
+ val row = recordReader.read()
+ rowsDecoded += 1
+ candidateByRelativeRow.get(relativeRow).foreach { candidate =>
+ if (verifyKeys && decodedRecordKey(row) != candidate.getRecordKey) {
+ keyMismatches += candidate
+ } else {
+ fillSyntheticFields(row, candidate)
+ rows += row
+ rowsMaterialized += 1
+ }
+ }
+ }
+ decodeMs += elapsedMs(decodeStartNs)
+ }
+ } finally {
+ reader.close()
+ }
+
+ val physicalAfter = stats.map(_.getBytesRead).getOrElse(physicalBefore)
+ val metrics = FetchMetrics(
+ rerankCandidates = candidates.size,
+ rowGroupsTotal = rowGroups.size,
+ rowGroupsSelected = rowGroupsSelected,
+ pagesInSelectedRowGroups = pagesInSelectedRowGroups,
+ pagesSelected = pagesSelected,
+ offsetIndexHits = offsetIndexHits,
+ offsetIndexMissing = offsetIndexMissing,
+ rangedGets = rangedGets,
+ pageBytesFetched = pageBytesFetched,
+ rowsDecoded = rowsDecoded,
+ rowsMaterialized = rowsMaterialized,
+ fetchKeyMismatch = keyMismatches.size,
+ physicalBytesRead = math.max(0L, physicalAfter - physicalBefore),
+ footerReads = if (footerResult.cacheHit) 0 else 1,
+ footerCacheHits = if (footerResult.cacheHit) 1 else 0,
+ footerMs = footerMs,
+ offsetIndexMs = offsetIndexMs,
+ fetchWaitMs = 0L,
+ decodeMs = decodeMs,
+ scoreMs = scoreMs,
+ elapsedMs = elapsedMs(startNs))
+ FetchResult(rows.iterator, keyMismatches.toSeq, metrics)
+ }
+
+ private def parquetRecordReader(
+ fileSchema: MessageType,
+ requestedParquetSchema: MessageType,
+ pageStore: PageReadStore): RecordReader[InternalRow] = {
+ val materializer = new InternalRowMaterializer(requestedParquetSchema, outputSchema, embeddingCol, distanceOrdinal)
+ new ColumnIOFactory().getColumnIO(requestedParquetSchema, fileSchema)
+ .getRecordReader(pageStore, materializer)
+ }
+
+ private def requestedSchema(fileSchema: MessageType): MessageType = {
+ val parquetFields =
+ if (vectorOnlyProjection) {
+ mutable.ArrayBuffer.empty[org.apache.parquet.schema.Type]
+ } else {
+ outputSchema.fields
+ .filterNot(field => isSyntheticOutputField(field.name))
+ .map(field => fieldByName(fileSchema, field.name))
+ .toBuffer[org.apache.parquet.schema.Type]
+ }
+ if ((verifyKeys || vectorOnlyProjection) && !parquetFields.exists(_.getName == HoodieRecord.RECORD_KEY_METADATA_FIELD)) {
+ parquetFields += fieldByName(fileSchema, HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ }
+ if (!parquetFields.exists(_.getName == embeddingCol)) {
+ parquetFields += fieldByName(fileSchema, embeddingCol)
+ }
+ new MessageType(fileSchema.getName, parquetFields.asJava)
+ }
+
+ private def isSyntheticOutputField(name: String): Boolean =
+ name == HoodieVectorSearchPlanBuilder.DISTANCE_COL ||
+ name == HoodieRecord.RECORD_KEY_METADATA_FIELD ||
+ name == HoodieRecord.PARTITION_PATH_METADATA_FIELD ||
+ name == HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL
+
+ private def fieldByName(schema: MessageType, name: String): org.apache.parquet.schema.Type =
+ schema.getFields.asScala.find(_.getName == name)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet schema is missing requested field '$name'."))
+
+ private def rowPosition(candidate: ScoredVectorPostingMatch): Long = {
+ val location = candidate.getLocation
+ if (location == null) {
+ throw new IllegalArgumentException(
+ s"Vector exact positional fetch requires a row-position locator for record ${candidate.getRecordKey}.")
+ }
+ location.getPosition
+ }
+
+ private def decodedRecordKey(row: InternalRow): String =
+ recordKeyOrdinal.flatMap(ordinal => Option(row.getUTF8String(ordinal))).map(_.toString).orNull
+
+ private def fillSyntheticFields(row: InternalRow, candidate: ScoredVectorPostingMatch): Unit = {
+ val liveLocation = Option(candidate.getLocation)
+ recordKeyOrdinal.foreach(row.update(_, UTF8String.fromString(candidate.getRecordKey)))
+ partitionPathOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getPartitionPath).orElse(Option(candidate.getPartitionPath)).getOrElse(""))))
+ fileGroupIdOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getFileId).orElse(Option(candidate.getFileGroupId)).getOrElse(""))))
+ }
+
+ private def fieldOrdinal(name: String): Option[Int] =
+ outputSchema.fields.indexWhere(_.name == name) match {
+ case -1 => None
+ case ordinal => Some(ordinal)
+ }
+
+ private def scoreVector(bytes: Array[Byte]): Double = scorer.scoreBytes(bytes)
+
+ private final class InternalRowMaterializer(
+ requestedParquetSchema: MessageType,
+ outputSchema: StructType,
+ embeddingCol: String,
+ distanceOrdinal: Int) extends RecordMaterializer[InternalRow] {
+
+ private val values = new Array[Any](outputSchema.length)
+ private var current: InternalRow = _
+ private var vectorBytes: Array[Byte] = _
+
+ private val root = new GroupConverter {
+ private val converters = requestedParquetSchema.getFields.asScala.map(converterFor).toIndexedSeq
+
+ override def getConverter(fieldIndex: Int): Converter =
+ converters(fieldIndex)
+
+ override def start(): Unit = {
+ var idx = 0
+ while (idx < values.length) {
+ values(idx) = null
+ idx += 1
+ }
+ vectorBytes = null
+ }
+
+ override def end(): Unit = {
+ if (vectorBytes == null) {
+ throw new IllegalStateException(s"Vector column '$embeddingCol' was not decoded for positional exact fetch.")
+ }
+ val scoreStartNs = System.nanoTime()
+ values(distanceOrdinal) = scoreVector(vectorBytes)
+ scoreMs += elapsedMs(scoreStartNs)
+ current = new GenericInternalRow(values.clone())
+ }
+ }
+
+ override def getCurrentRecord: InternalRow = current
+
+ override def getRootConverter: GroupConverter = root
+
+ private def converterFor(field: org.apache.parquet.schema.Type): PrimitiveConverter = {
+ val name = field.getName
+ new PrimitiveConverter {
+ override def addBinary(value: Binary): Unit = {
+ if (name == embeddingCol) {
+ vectorBytes = value.getBytes
+ } else {
+ setOutputValue(name, value.getBytes)
+ }
+ }
+
+ override def addInt(value: Int): Unit = setOutputValue(name, value)
+
+ override def addLong(value: Long): Unit = setOutputValue(name, value)
+
+ override def addFloat(value: Float): Unit = setOutputValue(name, value)
+
+ override def addDouble(value: Double): Unit = setOutputValue(name, value)
+
+ override def addBoolean(value: Boolean): Unit = setOutputValue(name, value)
+ }
+ }
+
+ private def setOutputValue(name: String, raw: Any): Unit = {
+ if (name != embeddingCol && name != HoodieVectorSearchPlanBuilder.DISTANCE_COL) {
+ val ordinal = outputSchema.fieldIndex(name)
+ values(ordinal) = outputSchema.fields(ordinal).dataType match {
+ case StringType => raw match {
+ case bytes: Array[Byte] => UTF8String.fromBytes(bytes)
+ case other => UTF8String.fromString(String.valueOf(other))
+ }
+ case BinaryType => raw.asInstanceOf[Array[Byte]]
+ case IntegerType => raw.asInstanceOf[Int]
+ case LongType => raw.asInstanceOf[Long]
+ case FloatType => raw.asInstanceOf[Float]
+ case DoubleType => raw.asInstanceOf[Double]
+ case BooleanType => raw.asInstanceOf[Boolean]
+ case ByteType => raw.asInstanceOf[Int].toByte
+ case ShortType => raw.asInstanceOf[Int].toShort
+ case other => throw new UnsupportedOperationException(
+ s"Unsupported exact positional fetch output type for field '$name': $other")
+ }
+ }
+ }
+ }
+}
+
+private[analysis] object PositionalParquetFetcher {
+ private val LOG = LoggerFactory.getLogger(getClass)
+ private val FOOTER_CACHE = new ConcurrentHashMap[String, ParquetMetadata]()
Review Comment:
🤖 non-blocking: these two static caches are process-lifetime and unbounded, keyed by path+length with no eviction when files are cleaned or clustered away, so a long-lived executor (STS/Kyuubi) accumulates a `ParquetMetadata` per base file ever touched; would a bounded LRU (or a per-query scope) make sense here?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorExactSearchPlanner.scala:
##########
@@ -0,0 +1,390 @@
+/*
+ * 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.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.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, 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}
+
+/** Exact reranking and file-slice resolution for the IVF/RaBitQ search plan. */
+private[analysis] object VectorExactSearchPlanner extends SparkAdapterSupport {
+ import HoodieVectorSearchPlanBuilder._
+ import IvfRaBitQMdtSearchAlgorithm.{isLogResidentCandidate, isPartitionedRecordIndex, toVectorMetric, VectorSearchPlanContext}
+
+ private val LOG = LoggerFactory.getLogger(getClass)
+
+ def buildRaBitQExactCandidateDf(
+ spark: SparkSession,
+ ctx: VectorSearchPlanContext,
+ outputSchema: StructType,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Float],
+ exactQueryVector: Array[Double],
+ k: Int): DataFrame = {
+ val planStartNs = System.nanoTime()
+ val storageConf = HadoopFSUtils.getStorageConf(spark.sessionState.newHadoopConf())
+ val cacheLoadStartNs = 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 topClusters = cache.findTopClusters(queryVector, numProbes, metric)
+ val topClusterSeq = topClusters.toSeq
+ val cacheLoadMs = elapsedMs(cacheLoadStartNs)
+
+ // Shard counts from cache — no extra MDT IO.
+ val shardCounts = cache.getShardCounts(topClusters)
+ val shardCountsScala = shardCounts.asScala.map { case (clusterId, count) =>
+ clusterId.intValue() -> count.intValue()
+ }.toMap
+ val postingPrefixCount = VectorIndexMdtSearchUtils.buildPostingPrefixes(cache.getGenerationId, shardCounts).size()
+ LOG.info(
+ s"[vector_search][plan] basePath=${ctx.basePath} indexPartition=${ctx.indexPartition} cacheLoadMs=$cacheLoadMs " +
+ s"generationId=${cache.getGenerationId} dimension=${cache.getDimension} numClusters=${cache.numClusters()} " +
+ s"numProbes=$numProbes refineFactor=${VectorIndexOptions.getRefineFactor(mergedOptions.asJava)} residualEncoding=$residualEncoding " +
+ s"topClusters=${summarizeInts(topClusterSeq)} shardCounts=${summarizeShardCounts(shardCountsScala)} " +
+ s"postingPrefixes=$postingPrefixCount")
+ if (shardCounts.isEmpty) {
+ sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+ } else {
+ val refineK = math.max(k, k * VectorIndexOptions.getRefineFactor(mergedOptions.asJava))
+ val emptyDf = sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+
+ 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, refineK)
+ val topCandidateMaterializeStartNs = System.nanoTime()
+ val materializedCandidates = try {
+ topCandidates.collectAsList().asScala.toSeq
+ } finally {
+ topCandidates.unpersistWithDependencies()
+ }
+ val topCandidateList = materializedCandidates
+ .filter(candidate => candidate.getFileGroupId != null && candidate.getFileGroupId.nonEmpty)
+ val candidatesWithoutLocator = materializedCandidates.size - topCandidateList.size
+ LOG.info(
+ s"[vector_search][stage][materialize_top_candidates] refineK=$refineK rerankCandidates=${materializedCandidates.size} " +
+ s"candidatesWithLocators=${topCandidateList.size} candidatesWithoutLocators=$candidatesWithoutLocator " +
+ s"distinctFileGroups=${topCandidateList.map(_.getFileGroupId).distinct.size} elapsedMs=${elapsedMs(topCandidateMaterializeStartNs)}")
+ if (candidatesWithoutLocator > 0) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search found $candidatesWithoutLocator rerank candidates without file-group locators. " +
+ s"Exact rerank requires per-candidate posting locators.")
+ }
+
+ // RFC-109 RLI finalist arbiter. Exact mode uses refreshed RLI locations for every live
+ // candidate: SERVE retains positional trust; STALE is resolved by key in the live base file.
+ val arbiterEnabled = VectorIndexOptions.isFinalistArbiterEnabled(mergedOptions.asJava)
+ val serveCandidates: Seq[ScoredVectorPostingMatch] =
+ if (arbiterEnabled) {
+ val arb = VectorIndexMdtSearchUtils.arbitrateMaterializedFinalists(
+ ctx.metadataTable, topCandidateList.asJava, isPartitionedRecordIndex(ctx))
+ val staleFailPolicy = VectorIndexOptions.isStaleLocatorPolicyFail(mergedOptions.asJava)
+ LOG.info(
+ s"[vector_search][exact][arbiter] arbiterExclusions{stale=${arb.staleCount}, deleted=${arb.deletedCount}} " +
+ s"serve=${arb.serve.size} stalePolicy=${if (staleFailPolicy) "fail" else "fallback"}")
+ if (staleFailPolicy && !arb.stale.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search rejected ${arb.staleCount} stale finalist(s) because " +
+ "vector.stale.locator.policy=FAIL.")
+ }
+ arb.serve.asScala.toSeq ++ arb.stale.asScala.toSeq
+ } else {
+ topCandidateList
+ }
+
+ if (serveCandidates.isEmpty) {
+ emptyDf
+ } else {
+ val candidatesByPartition = serveCandidates
+ .groupBy(candidate => Option(candidate.getLocation).map(_.getPartitionPath).getOrElse(""))
+ .map { case (partitionPath, candidates) =>
+ partitionPath -> candidates.flatMap(candidate => Option(candidate.getLocation).map(_.getFileId)).toSet
+ }
+ val fileSliceMap = preResolveFileSlices(
+ ctx.metadataTable, ctx.metaClient, ctx.timeline, ctx.latestInstantTime, candidatesByPartition)
+ if (fileSliceMap.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search resolved zero file slices for ${serveCandidates.size} rerank candidates " +
+ s"across ${candidatesByPartition.values.map(_.size).sum} file groups. Candidate locators are present, so this indicates " +
+ s"a file-slice resolution bug or stale vector index.")
+ }
+ val bFileSliceMap = spark.sparkContext.broadcast(fileSliceMap)
+ // MOR log-resident candidates carry their raw embedding in a log block, not the base Parquet
+ // file, so the positional fetcher cannot read them. Split them out and materialize them via a
+ // merged (base + log) read; base-resident candidates keep the fast positional path.
+ val (logResidentCandidates, baseCandidates) = serveCandidates.partition { candidate =>
+ bFileSliceMap.value.get(candidate.getLocation.getFileId).exists { slice =>
+ isLogResidentCandidate(candidate, slice.getBaseInstantTime, slice.getLogFiles.findAny().isPresent)
+ }
+ }
+ val exactFetchParallelism = math.max(
+ 1,
+ math.min(refineK, math.max(1, spark.sparkContext.defaultParallelism * 4)))
+ val groupedCandidates = baseCandidates.groupBy(_.getLocation.getFileId).toSeq.sortBy(_._1)
+ val targetBatchesPerGroup = math.max(1, exactFetchParallelism / math.max(1, groupedCandidates.size))
+ val fetchBatches = groupedCandidates.flatMap { case (fgId, candidates) =>
+ val sorted = candidates.sortBy(candidate => candidate.getLocation.getPosition)
+ val batchSize = math.max(1, math.ceil(sorted.size.toDouble / targetBatchesPerGroup).toInt)
+ sorted.grouped(batchSize).map(batch => (fgId, batch)).toSeq
+ }
+ val byFetchBatch: RDD[(String, Seq[ScoredVectorPostingMatch])] =
+ if (fetchBatches.isEmpty) {
+ spark.sparkContext.emptyRDD[(String, Seq[ScoredVectorPostingMatch])]
+ } else {
+ spark.sparkContext.parallelize(fetchBatches, math.min(exactFetchParallelism, fetchBatches.size))
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read_plan] rerankCandidates=${serveCandidates.size} " +
+ s"candidateGroups=${groupedCandidates.size} resolvedSlices=${fileSliceMap.size} " +
+ s"fileGroupPartitions=${byFetchBatch.getNumPartitions} exactFetchParallelism=$exactFetchParallelism")
+
+ val vectorOnlyProjection = mergedOptions.get(VECTOR_ONLY_PROJECTION_OPT)
+ .orElse(spark.conf.getOption(VECTOR_ONLY_PROJECTION_SPARK_CONF))
+ .exists(_.equalsIgnoreCase("true"))
+ val fetchSkippedLogResidentAcc: LongAccumulator =
+ spark.sparkContext.longAccumulator("vector_fetch_skipped_log_resident")
+ val fetchKeyMismatchAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_fetch_key_mismatch")
+ val rowsMaterializedAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_rows_materialized")
+ val exactRows: RDD[InternalRow] = byFetchBatch.mapPartitions { batches =>
+ val startNs = System.nanoTime()
+ val rows = ArrayBuffer.empty[InternalRow]
+ var batchCount = 0
+ var resolvedSlices = 0
+ var missingSlices = 0
+ var candidateKeys = 0L
+ var baseFiles = 0
+ var logFiles = 0L
+ var totalSliceBytes = 0L
+ var staleCandidates = 0L
+ var fetchSkippedLogResident = 0L
+ var fetchKeyMismatch = 0L
+ var keyLookupRowsDecoded = 0L
+ var metrics = PositionalParquetFetcher.FetchMetrics()
+ val fetcher = new PositionalParquetFetcher(
+ storageConf.newInstance().unwrap(),
+ embeddingCol,
+ vectorSchema,
+ exactQueryVector,
+ ctx.metric,
+ outputSchema,
+ mergedOptions.getOrElse("vector.exact_fetch.offset_index_missing_policy", "fallback").equalsIgnoreCase("fail"),
+ vectorOnlyProjection,
+ VectorIndexOptions.shouldVerifyFetchKeys(mergedOptions.asJava))
+ val keyLocator = new ParquetRecordKeyLocator(storageConf.newInstance().unwrap())
+ while (batches.hasNext) {
+ val (fgId, candidates) = batches.next()
+ batchCount += 1
+ val fileSliceOpt = bFileSliceMap.value.get(fgId)
+ if (fileSliceOpt.isEmpty) {
+ missingSlices += 1
+ throw new HoodieAnalysisException(
+ s"Vector exact search could not resolve file group '$fgId' for ${candidates.size} rerank candidates. " +
+ s"No record-level-index fallback is active in this path.")
+ } else {
+ resolvedSlices += 1
+ val fileSlice = fileSliceOpt.get
+ val logFileCount = fileSlice.getLogFiles.count()
+ val currentBaseInstantTime = fileSlice.getBaseInstantTime
+ val logResident = candidates.filter(candidate =>
+ isLogResidentCandidate(candidate, currentBaseInstantTime, logFileCount > 0))
+ if (logResident.nonEmpty) {
+ fetchSkippedLogResident += logResident.size
+ fetchSkippedLogResidentAcc.add(logResident.size)
+ LOG.warn(
+ s"Vector exact fetch skipped ${logResident.size} log-resident candidate(s) in file group '$fgId'; " +
+ "they remain a result shortfall and the retained candidate pool will refill where candidates remain")
+ }
+ val liveCandidates = candidates.filterNot(logResident.toSet)
+ staleCandidates += liveCandidates.count(candidate =>
+ candidate.getArbiterDecision == VectorIndexArbiter.Decision.STALE)
+ if (liveCandidates.nonEmpty) {
+ val positionalCandidates = liveCandidates.filter { candidate =>
+ val location = candidate.getLocation
+ location != null && location.getPosition >= 0 &&
+ candidate.getArbiterDecision != VectorIndexArbiter.Decision.STALE &&
+ (location.getInstantTime == null || location.getInstantTime == currentBaseInstantTime)
+ }
+ val keyCandidates = liveCandidates.filterNot(positionalCandidates.toSet)
+ val recordKeys = liveCandidates.map(_.getRecordKey).toSet
+ candidateKeys += recordKeys.size
+ if (fileSlice.getBaseFile.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact key fallback requires a base Parquet file for file group '$fgId'; " +
+ "the remaining candidates are not classified as log-resident.")
+ }
+ baseFiles += 1
+ logFiles += logFileCount
+ totalSliceBytes += fileSlice.getTotalFileSize
+ val baseFilePath = fileSlice.getBaseFile.get.getPath
+ val locatedByKey = keyLocator.locate(baseFilePath, keyCandidates)
+ keyLookupRowsDecoded += locatedByKey.metrics.rowsDecoded
+ if (locatedByKey.metrics.located != keyCandidates.size) {
+ throw new HoodieAnalysisException(
+ s"Vector exact key fallback located ${locatedByKey.metrics.located} of ${keyCandidates.size} " +
+ s"candidate key(s) in the live base file for file group '$fgId'.")
+ }
+ val firstFetch = fetcher.fetch(baseFilePath, positionalCandidates ++ locatedByKey.candidates)
+ rows ++= firstFetch.rows
+ metrics = metrics.add(firstFetch.metrics)
+ rowsMaterializedAcc.add(firstFetch.metrics.rowsMaterialized)
+ if (firstFetch.retryByKey.nonEmpty) {
+ fetchKeyMismatch += firstFetch.retryByKey.size
+ fetchKeyMismatchAcc.add(firstFetch.retryByKey.size)
+ val retries = keyLocator.locate(baseFilePath, firstFetch.retryByKey)
+ keyLookupRowsDecoded += retries.metrics.rowsDecoded
+ if (retries.metrics.located != firstFetch.retryByKey.size) {
+ throw new HoodieAnalysisException(
+ s"Vector exact fetch could not relocate ${firstFetch.retryByKey.size - retries.metrics.located} " +
+ s"key mismatch(es) in file group '$fgId'.")
+ }
+ val retryFetch = fetcher.fetch(baseFilePath, retries.candidates)
+ if (retryFetch.retryByKey.nonEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact fetch could not self-heal ${retryFetch.retryByKey.size} key mismatch(es) " +
+ s"in file group '$fgId' after key-based relocation.")
+ }
+ rows ++= retryFetch.rows
+ metrics = metrics.add(retryFetch.metrics)
+ rowsMaterializedAcc.add(retryFetch.metrics.rowsMaterialized)
+ }
+ }
+ }
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read] rerankCandidates=$candidateKeys candidateFileCount=$resolvedSlices " +
+ s"fetchBatches=$batchCount resolvedSlices=$resolvedSlices missingSlices=$missingSlices " +
+ s"rowsDecoded=${metrics.rowsDecoded} rowsMaterialized=${metrics.rowsMaterialized} " +
+ s"keyLookupRowsDecoded=$keyLookupRowsDecoded fetchKeyMismatch=$fetchKeyMismatch " +
+ s"fetchSkippedLogResident=$fetchSkippedLogResident " +
+ s"rowGroupsTotal=${metrics.rowGroupsTotal} rowGroupsSelected=${metrics.rowGroupsSelected} " +
+ s"pagesInSelectedRowGroups=${metrics.pagesInSelectedRowGroups} pagesSelected=${metrics.pagesSelected} " +
+ s"offsetIndexHits=${metrics.offsetIndexHits} offsetIndexMissing=${metrics.offsetIndexMissing} " +
+ s"rangedGets=${metrics.rangedGets} pageBytesFetched=${metrics.pageBytesFetched} " +
+ s"physicalBytesRead=${metrics.physicalBytesRead} footerReads=${metrics.footerReads} footerCacheHits=${metrics.footerCacheHits} " +
+ s"fetchKeyMismatchMetric=${metrics.fetchKeyMismatch} staleCandidates=$staleCandidates " +
+ s"baseFiles=$baseFiles logFiles=$logFiles " +
+ s"sliceBytes=$totalSliceBytes footerMs=${metrics.footerMs} offsetIndexMs=${metrics.offsetIndexMs} " +
+ s"fetchWaitMs=${metrics.fetchWaitMs} decodeMs=${metrics.decodeMs} scoreMs=${metrics.scoreMs} elapsedMs=${elapsedMs(startNs)}")
+ rows.iterator
+ }
+ val logResidentRows: Seq[InternalRow] =
+ if (logResidentCandidates.isEmpty) {
+ Seq.empty
+ } else {
+ val dataSchema = new TableSchemaResolver(ctx.metaClient).getTableSchema(true)
+ val scorer = new VectorExactScorer(vectorSchema, exactQueryVector, ctx.metric)
+ val logFetcher = new LogResidentVectorFetcher(
+ storageConf, ctx.metaClient, dataSchema, ctx.latestInstantTime,
+ embeddingCol, outputSchema, scorer)
+ val fetched = logResidentCandidates
+ .groupBy(candidate =>
+ (Option(candidate.getLocation.getPartitionPath).getOrElse(""), candidate.getLocation.getFileId))
+ .flatMap { case ((partitionPath, fgId), candidates) =>
+ logFetcher.fetchSlice(partitionPath, fileSliceMap(fgId), candidates.map(_.getRecordKey).toSet)
+ }.toSeq
+ LOG.info(
+ s"[vector_search][exact][log_resident] candidates=${logResidentCandidates.size} " +
+ s"materialized=${fetched.size} " +
+ s"fileGroups=${logResidentCandidates.map(_.getLocation.getFileId).distinct.size}")
+ fetched
+ }
+ val combinedRows: RDD[InternalRow] =
+ if (logResidentRows.isEmpty) exactRows
+ else exactRows.union(spark.sparkContext.parallelize(logResidentRows))
+ LOG.info(s"[vector_search][plan] exact candidate DF planned in ${elapsedMs(planStartNs)} ms")
+ sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(spark, combinedRows, outputSchema)
+ }
+ }
+ }
+
+ private def preResolveFileSlices(
+ metadataTable: HoodieTableMetadata,
+ metaClient: HoodieTableMetaClient,
+ timeline: HoodieTimeline,
+ latestInstantTime: String,
+ candidatesByPartition: Map[String, Set[String]]): Map[String, FileSlice] = {
+ val startNs = System.nanoTime()
+ val fsView = new HoodieTableFileSystemView(metadataTable, metaClient, timeline)
+ try {
+ val result = scala.collection.mutable.Map.empty[String, FileSlice]
+ var baseFiles = 0
+ var logFiles = 0L
+ var totalSliceBytes = 0L
+ candidatesByPartition.foreach { case (partitionPath, candidateFgIds) =>
+ val sliceIter = fsView.getLatestFileSlicesBeforeOrOn(partitionPath, latestInstantTime, true).iterator()
+ while (sliceIter.hasNext) {
+ val slice = sliceIter.next()
+ val fgId = slice.getFileGroupId.getFileId
+ if (candidateFgIds.contains(fgId)) {
+ result(fgId) = slice
Review Comment:
🤖 non-blocking: the slice map is keyed by file id alone, but NBCC bucket-index tables reuse the same fixed file id (`{bucket}-0000-0000-0000-000000000000-0`) in every partition, so candidates from two partitions in one bucket resolve to whichever partition's slice was visited last (wrong base file for the positional/key fetch, silent misses on the log-resident path); keying by `HoodieFileGroupId` here and in the lookups at L168/L335 would avoid that.
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
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()));
+ }
+ return finalists.map(candidate ->
+ arbitrateCandidate(candidate, currentLocations.get(candidate.getRecordKey())));
+ }
+
+ private static ScoredVectorPostingMatch arbitrateCandidate(
+ ScoredVectorPostingMatch candidate,
+ HoodieRecordGlobalLocation current) {
+ VectorIndexArbiter.Decision decision = VectorIndexArbiter.classify(
+ candidate.getPartitionPath(),
+ candidate.getFileGroupId(),
+ candidate.getBaseInstantTime(),
+ current);
+ HoodieRecordGlobalLocation resolved;
+ switch (decision) {
+ case SERVE:
+ resolved = candidate.getLocation() != null ? candidate.getLocation() : current;
+ break;
+ case STALE:
+ resolved = current;
+ break;
+ case DELETED:
+ default:
+ resolved = null;
+ break;
+ }
+ return candidate.withArbiterVerdict(decision, resolved);
+ }
+
+ /**
+ * Driver-side finalist arbiter core: classify each already-materialized finalist against a
+ * pre-resolved map of current RLI locations (record key -> location, or absent for an RLI miss).
+ * Pure and Spark-free so it is directly unit-testable; the {@link HoodieTableMetadata} overload
+ * performs the batched RLI lookup and delegates here.
+ *
+ * <p>Implements the RFC-109 arbiter output contract (see {@link VectorIndexArbiter}):
+ * hit+match -> SERVE (positional trust preserved via the posting's own location when present),
+ * hit+differ -> STALE (resolved to the live RLI location), miss -> DELETED (dropped).
+ */
+ static VectorIndexMdtSearchUtils.ArbitrationResult arbitrateMaterializedFinalists(
+ List<ScoredVectorPostingMatch> finalists,
+ Map<String, HoodieRecordGlobalLocation> currentLocations) {
+ List<ScoredVectorPostingMatch> serve = new ArrayList<>();
+ List<ScoredVectorPostingMatch> stale = new ArrayList<>();
+ long deleted = 0L;
+ for (ScoredVectorPostingMatch candidate : finalists) {
+ HoodieRecordGlobalLocation current = currentLocations.get(candidate.getRecordKey());
+ VectorIndexArbiter.Decision decision = VectorIndexArbiter.classify(
+ candidate.getPartitionPath(),
+ candidate.getFileGroupId(),
+ candidate.getBaseInstantTime(),
+ current);
Review Comment:
🤖 nit: this switch duplicates `arbitrateCandidate`'s resolved-location rules — could the loop call `arbitrateCandidate` and then bucket by the returned verdict, so the SERVE/STALE/DELETED semantics are defined once?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorExactSearchPlanner.scala:
##########
@@ -0,0 +1,390 @@
+/*
+ * 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.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.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, 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}
+
+/** Exact reranking and file-slice resolution for the IVF/RaBitQ search plan. */
+private[analysis] object VectorExactSearchPlanner extends SparkAdapterSupport {
+ import HoodieVectorSearchPlanBuilder._
+ import IvfRaBitQMdtSearchAlgorithm.{isLogResidentCandidate, isPartitionedRecordIndex, toVectorMetric, VectorSearchPlanContext}
+
+ private val LOG = LoggerFactory.getLogger(getClass)
+
+ def buildRaBitQExactCandidateDf(
+ spark: SparkSession,
+ ctx: VectorSearchPlanContext,
+ outputSchema: StructType,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Float],
+ exactQueryVector: Array[Double],
+ k: Int): DataFrame = {
+ val planStartNs = System.nanoTime()
+ val storageConf = HadoopFSUtils.getStorageConf(spark.sessionState.newHadoopConf())
+ val cacheLoadStartNs = 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 topClusters = cache.findTopClusters(queryVector, numProbes, metric)
+ val topClusterSeq = topClusters.toSeq
+ val cacheLoadMs = elapsedMs(cacheLoadStartNs)
+
+ // Shard counts from cache — no extra MDT IO.
+ val shardCounts = cache.getShardCounts(topClusters)
+ val shardCountsScala = shardCounts.asScala.map { case (clusterId, count) =>
+ clusterId.intValue() -> count.intValue()
+ }.toMap
+ val postingPrefixCount = VectorIndexMdtSearchUtils.buildPostingPrefixes(cache.getGenerationId, shardCounts).size()
+ LOG.info(
+ s"[vector_search][plan] basePath=${ctx.basePath} indexPartition=${ctx.indexPartition} cacheLoadMs=$cacheLoadMs " +
+ s"generationId=${cache.getGenerationId} dimension=${cache.getDimension} numClusters=${cache.numClusters()} " +
+ s"numProbes=$numProbes refineFactor=${VectorIndexOptions.getRefineFactor(mergedOptions.asJava)} residualEncoding=$residualEncoding " +
+ s"topClusters=${summarizeInts(topClusterSeq)} shardCounts=${summarizeShardCounts(shardCountsScala)} " +
+ s"postingPrefixes=$postingPrefixCount")
+ if (shardCounts.isEmpty) {
+ sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+ } else {
+ val refineK = math.max(k, k * VectorIndexOptions.getRefineFactor(mergedOptions.asJava))
+ val emptyDf = sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+
+ 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, refineK)
+ val topCandidateMaterializeStartNs = System.nanoTime()
+ val materializedCandidates = try {
+ topCandidates.collectAsList().asScala.toSeq
+ } finally {
+ topCandidates.unpersistWithDependencies()
+ }
+ val topCandidateList = materializedCandidates
+ .filter(candidate => candidate.getFileGroupId != null && candidate.getFileGroupId.nonEmpty)
+ val candidatesWithoutLocator = materializedCandidates.size - topCandidateList.size
+ LOG.info(
+ s"[vector_search][stage][materialize_top_candidates] refineK=$refineK rerankCandidates=${materializedCandidates.size} " +
+ s"candidatesWithLocators=${topCandidateList.size} candidatesWithoutLocators=$candidatesWithoutLocator " +
+ s"distinctFileGroups=${topCandidateList.map(_.getFileGroupId).distinct.size} elapsedMs=${elapsedMs(topCandidateMaterializeStartNs)}")
+ if (candidatesWithoutLocator > 0) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search found $candidatesWithoutLocator rerank candidates without file-group locators. " +
+ s"Exact rerank requires per-candidate posting locators.")
+ }
+
+ // RFC-109 RLI finalist arbiter. Exact mode uses refreshed RLI locations for every live
+ // candidate: SERVE retains positional trust; STALE is resolved by key in the live base file.
+ val arbiterEnabled = VectorIndexOptions.isFinalistArbiterEnabled(mergedOptions.asJava)
+ val serveCandidates: Seq[ScoredVectorPostingMatch] =
+ if (arbiterEnabled) {
+ val arb = VectorIndexMdtSearchUtils.arbitrateMaterializedFinalists(
+ ctx.metadataTable, topCandidateList.asJava, isPartitionedRecordIndex(ctx))
+ val staleFailPolicy = VectorIndexOptions.isStaleLocatorPolicyFail(mergedOptions.asJava)
+ LOG.info(
+ s"[vector_search][exact][arbiter] arbiterExclusions{stale=${arb.staleCount}, deleted=${arb.deletedCount}} " +
+ s"serve=${arb.serve.size} stalePolicy=${if (staleFailPolicy) "fail" else "fallback"}")
+ if (staleFailPolicy && !arb.stale.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search rejected ${arb.staleCount} stale finalist(s) because " +
+ "vector.stale.locator.policy=FAIL.")
+ }
+ arb.serve.asScala.toSeq ++ arb.stale.asScala.toSeq
+ } else {
+ topCandidateList
+ }
+
+ if (serveCandidates.isEmpty) {
+ emptyDf
+ } else {
+ val candidatesByPartition = serveCandidates
+ .groupBy(candidate => Option(candidate.getLocation).map(_.getPartitionPath).getOrElse(""))
+ .map { case (partitionPath, candidates) =>
+ partitionPath -> candidates.flatMap(candidate => Option(candidate.getLocation).map(_.getFileId)).toSet
+ }
+ val fileSliceMap = preResolveFileSlices(
+ ctx.metadataTable, ctx.metaClient, ctx.timeline, ctx.latestInstantTime, candidatesByPartition)
+ if (fileSliceMap.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search resolved zero file slices for ${serveCandidates.size} rerank candidates " +
+ s"across ${candidatesByPartition.values.map(_.size).sum} file groups. Candidate locators are present, so this indicates " +
+ s"a file-slice resolution bug or stale vector index.")
+ }
+ val bFileSliceMap = spark.sparkContext.broadcast(fileSliceMap)
+ // MOR log-resident candidates carry their raw embedding in a log block, not the base Parquet
+ // file, so the positional fetcher cannot read them. Split them out and materialize them via a
+ // merged (base + log) read; base-resident candidates keep the fast positional path.
+ val (logResidentCandidates, baseCandidates) = serveCandidates.partition { candidate =>
+ bFileSliceMap.value.get(candidate.getLocation.getFileId).exists { slice =>
+ isLogResidentCandidate(candidate, slice.getBaseInstantTime, slice.getLogFiles.findAny().isPresent)
+ }
+ }
+ val exactFetchParallelism = math.max(
+ 1,
+ math.min(refineK, math.max(1, spark.sparkContext.defaultParallelism * 4)))
+ val groupedCandidates = baseCandidates.groupBy(_.getLocation.getFileId).toSeq.sortBy(_._1)
+ val targetBatchesPerGroup = math.max(1, exactFetchParallelism / math.max(1, groupedCandidates.size))
+ val fetchBatches = groupedCandidates.flatMap { case (fgId, candidates) =>
+ val sorted = candidates.sortBy(candidate => candidate.getLocation.getPosition)
+ val batchSize = math.max(1, math.ceil(sorted.size.toDouble / targetBatchesPerGroup).toInt)
+ sorted.grouped(batchSize).map(batch => (fgId, batch)).toSeq
+ }
+ val byFetchBatch: RDD[(String, Seq[ScoredVectorPostingMatch])] =
+ if (fetchBatches.isEmpty) {
+ spark.sparkContext.emptyRDD[(String, Seq[ScoredVectorPostingMatch])]
+ } else {
+ spark.sparkContext.parallelize(fetchBatches, math.min(exactFetchParallelism, fetchBatches.size))
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read_plan] rerankCandidates=${serveCandidates.size} " +
+ s"candidateGroups=${groupedCandidates.size} resolvedSlices=${fileSliceMap.size} " +
+ s"fileGroupPartitions=${byFetchBatch.getNumPartitions} exactFetchParallelism=$exactFetchParallelism")
+
+ val vectorOnlyProjection = mergedOptions.get(VECTOR_ONLY_PROJECTION_OPT)
+ .orElse(spark.conf.getOption(VECTOR_ONLY_PROJECTION_SPARK_CONF))
+ .exists(_.equalsIgnoreCase("true"))
+ val fetchSkippedLogResidentAcc: LongAccumulator =
+ spark.sparkContext.longAccumulator("vector_fetch_skipped_log_resident")
+ val fetchKeyMismatchAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_fetch_key_mismatch")
+ val rowsMaterializedAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_rows_materialized")
+ val exactRows: RDD[InternalRow] = byFetchBatch.mapPartitions { batches =>
+ val startNs = System.nanoTime()
+ val rows = ArrayBuffer.empty[InternalRow]
+ var batchCount = 0
+ var resolvedSlices = 0
+ var missingSlices = 0
+ var candidateKeys = 0L
+ var baseFiles = 0
+ var logFiles = 0L
+ var totalSliceBytes = 0L
+ var staleCandidates = 0L
+ var fetchSkippedLogResident = 0L
+ var fetchKeyMismatch = 0L
+ var keyLookupRowsDecoded = 0L
+ var metrics = PositionalParquetFetcher.FetchMetrics()
+ val fetcher = new PositionalParquetFetcher(
+ storageConf.newInstance().unwrap(),
+ embeddingCol,
+ vectorSchema,
+ exactQueryVector,
+ ctx.metric,
Review Comment:
🤖 non-blocking: referencing `ctx.metric` inside the `mapPartitions` closure captures the whole `VectorSearchPlanContext` (meta client, metadata table with its file-system view, and the `float[][]` centroid cache) into the exact stage's task binary on every query; could you hoist it into a local above the closure, like `mergedOptions` already is?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/PositionalParquetFetcher.scala:
##########
@@ -0,0 +1,587 @@
+/*
+ * 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.spark.sql.hudi.analysis
+
+import org.apache.hudi.common.index.vector.{ScoredVectorPostingMatch, VectorIndexMdtSearchUtils, VectorPostingMatch}
+import org.apache.hudi.common.model.HoodieRecord
+import org.apache.hudi.common.schema.HoodieSchema
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileSystem, Path}
+import org.apache.parquet.column.page.PageReadStore
+import org.apache.parquet.hadoop.ParquetFileReader
+import org.apache.parquet.hadoop.metadata.{BlockMetaData, ColumnChunkMetaData, ParquetMetadata}
+import org.apache.parquet.internal.column.columnindex.OffsetIndex
+import org.apache.parquet.internal.filter2.columnindex.RowRanges
+import org.apache.parquet.io.{ColumnIOFactory, RecordReader}
+import org.apache.parquet.io.api.{Binary, Converter, GroupConverter, PrimitiveConverter, RecordMaterializer}
+import org.apache.parquet.schema.MessageType
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction.DistanceMetric
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+import org.slf4j.LoggerFactory
+
+import java.util.PrimitiveIterator
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+/**
+ * Fetches exact vector-search rerank candidates from Parquet by physical row position.
+ *
+ * The fetcher uses Parquet page indexes to turn candidate row positions into row ranges, so
+ * the reader can skip non-candidate pages instead of scanning whole Hudi file slices.
+ */
+private[analysis] final class PositionalParquetFetcher(
+ conf: Configuration,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Double],
+ metric: DistanceMetric.Value,
+ outputSchema: StructType,
+ failOnMissingOffsetIndex: Boolean,
+ vectorOnlyProjection: Boolean,
+ verifyKeys: Boolean) extends Serializable {
+
+ import PositionalParquetFetcher._
+
+ private val distanceOrdinal = outputSchema.fieldIndex(HoodieVectorSearchPlanBuilder.DISTANCE_COL)
+ private val recordKeyOrdinal = fieldOrdinal(HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ private val partitionPathOrdinal = fieldOrdinal(HoodieRecord.PARTITION_PATH_METADATA_FIELD)
+ private val fileGroupIdOrdinal = fieldOrdinal(HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL)
+ private val scorer = new VectorExactScorer(vectorSchema, queryVector, metric)
+ private var scoreMs: Long = 0L
+
+ def fetch(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch]): FetchResult = {
+ val startNs = System.nanoTime()
+ if (candidates.isEmpty) {
+ FetchResult(Iterator.empty, Seq.empty, FetchMetrics(elapsedMs = elapsedMs(startNs)))
+ } else {
+ fetchNonEmpty(filePath, candidates, startNs)
+ }
+ }
+
+ private def fetchNonEmpty(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch],
+ startNs: Long): FetchResult = {
+ val path = new Path(filePath)
+ configureRandomRead(conf)
+ val fs = path.getFileSystem(conf)
+ val fileLen = fs.getFileStatus(path).getLen
+ val fileCacheKey = s"$filePath#$fileLen"
+ val stats = Option(FileSystem.getStatistics(path.toUri.getScheme, fs.getClass))
+ val physicalBefore = stats.map(_.getBytesRead).getOrElse(0L)
+
+ val footerStartNs = System.nanoTime()
+ val footerResult = footer(conf, path, fileCacheKey)
+ val footerMs = elapsedMs(footerStartNs)
+ val metadata = footerResult.metadata
+ val fileSchema = metadata.getFileMetaData.getSchema
+ val requestedParquetSchema = requestedSchema(fileSchema)
+
+ val rowGroups = metadata.getBlocks.asScala.toIndexedSeq
+ val rowGroupStarts = cumulativeRowGroupStarts(rowGroups)
+ val byRowGroup: Map[Int, Seq[ScoredVectorPostingMatch]] =
+ candidates.groupBy(candidate => rowGroupOrdinal(rowGroupStarts, rowGroups, rowPosition(candidate)))
+
+ val rows = mutable.ArrayBuffer.empty[InternalRow]
+ val keyMismatches = mutable.ArrayBuffer.empty[ScoredVectorPostingMatch]
+ var rowGroupsSelected = 0
+ var pagesInSelectedRowGroups = 0L
+ var pagesSelected = 0L
+ var offsetIndexHits = 0L
+ var offsetIndexMissing = 0L
+ var pageBytesFetched = 0L
+ var rangedGets = 0L
+ var rowsDecoded = 0L
+ var rowsMaterialized = 0L
+ var offsetIndexMs = 0L
+ var decodeMs = 0L
+ scoreMs = 0L
+
+ val reader = new ParquetFileReader(conf, path, metadata)
+ try {
+ reader.setRequestedSchema(requestedParquetSchema)
+ byRowGroup.toSeq.sortBy(_._1).foreach { case (rgOrdinal, rgCandidates) =>
+ val block = rowGroups(rgOrdinal)
+ val rowGroupStart = rowGroupStarts(rgOrdinal)
+ val embeddingChunk = findColumn(block, embeddingCol)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet file '$filePath' is missing vector column '$embeddingCol'."))
+
+ val offsetStartNs = System.nanoTime()
+ val embeddingOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, embeddingChunk)
+ offsetIndexMs += elapsedMs(offsetStartNs)
+ val relativeRows = rgCandidates.map(candidate => rowPosition(candidate) - rowGroupStart).sorted
+ val candidateRows = relativeRows.toSet
+ val candidateByRelativeRow = rgCandidates
+ .map(candidate => rowPosition(candidate) - rowGroupStart -> candidate)
+ .toMap
+
+ if (embeddingOffsetIndex == null && failOnMissingOffsetIndex) {
+ offsetIndexMissing += 1
+ throw new IllegalStateException(
+ s"Parquet file '$filePath' has no OffsetIndex for vector column '$embeddingCol'; positional exact fetch cannot honor pages.")
+ }
+
+ rowGroupsSelected += 1
+ val (pageStore, rowRanges) = if (embeddingOffsetIndex == null) {
+ offsetIndexMissing += 1
+ LOG.warn(
+ "Parquet file {} rowGroup={} has no OffsetIndex for vector column {}; falling back to row-group fetch.",
+ filePath,
+ Int.box(rgOrdinal),
+ embeddingCol)
+ pagesInSelectedRowGroups += 0L
+ pagesSelected += 0L
+ rangedGets += 1L
+ pageBytesFetched += requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .map(_.getTotalSize)
+ .sum
+ (reader.readRowGroup(rgOrdinal), RowRanges.createSingle(block.getRowCount))
+ } else {
+ offsetIndexHits += 1
+ pagesInSelectedRowGroups += embeddingOffsetIndex.getPageCount
+ val selectedVectorPages = selectedPages(embeddingOffsetIndex, block.getRowCount, relativeRows)
+ pagesSelected += selectedVectorPages.size
+ val selectedRanges = requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .flatMap { chunk =>
+ val oiStartNs = System.nanoTime()
+ val chunkOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, chunk)
+ offsetIndexMs += elapsedMs(oiStartNs)
+ if (chunkOffsetIndex == null) {
+ offsetIndexMissing += 1
+ Seq.empty
+ } else {
+ offsetIndexHits += 1
+ selectedPages(chunkOffsetIndex, block.getRowCount, relativeRows).map { pageOrdinal =>
+ val start = chunkOffsetIndex.getOffset(pageOrdinal)
+ PageRange(start, start + chunkOffsetIndex.getCompressedPageSize(pageOrdinal))
+ }
+ }
+ }
+ val mergedRanges = mergeRanges(selectedRanges.sortBy(_.start).toSeq)
+ rangedGets += mergedRanges.size
+ pageBytesFetched += mergedRanges.iterator.map(range => range.end - range.start).sum
+ val ranges = RowRanges.create(block.getRowCount, primitiveIntIterator(selectedVectorPages), embeddingOffsetIndex)
+ (reader.readFilteredRowGroup(rgOrdinal, ranges), ranges)
+ }
+
+ val decodeStartNs = System.nanoTime()
+ val recordReader = parquetRecordReader(fileSchema, requestedParquetSchema, pageStore)
+ val rowIterator = rowRanges.iterator()
+ while (rowIterator.hasNext) {
+ val relativeRow = rowIterator.nextLong()
+ val row = recordReader.read()
+ rowsDecoded += 1
+ candidateByRelativeRow.get(relativeRow).foreach { candidate =>
+ if (verifyKeys && decodedRecordKey(row) != candidate.getRecordKey) {
+ keyMismatches += candidate
+ } else {
+ fillSyntheticFields(row, candidate)
+ rows += row
+ rowsMaterialized += 1
+ }
+ }
+ }
+ decodeMs += elapsedMs(decodeStartNs)
+ }
+ } finally {
+ reader.close()
+ }
+
+ val physicalAfter = stats.map(_.getBytesRead).getOrElse(physicalBefore)
+ val metrics = FetchMetrics(
+ rerankCandidates = candidates.size,
+ rowGroupsTotal = rowGroups.size,
+ rowGroupsSelected = rowGroupsSelected,
+ pagesInSelectedRowGroups = pagesInSelectedRowGroups,
+ pagesSelected = pagesSelected,
+ offsetIndexHits = offsetIndexHits,
+ offsetIndexMissing = offsetIndexMissing,
+ rangedGets = rangedGets,
+ pageBytesFetched = pageBytesFetched,
+ rowsDecoded = rowsDecoded,
+ rowsMaterialized = rowsMaterialized,
+ fetchKeyMismatch = keyMismatches.size,
+ physicalBytesRead = math.max(0L, physicalAfter - physicalBefore),
+ footerReads = if (footerResult.cacheHit) 0 else 1,
+ footerCacheHits = if (footerResult.cacheHit) 1 else 0,
+ footerMs = footerMs,
+ offsetIndexMs = offsetIndexMs,
+ fetchWaitMs = 0L,
+ decodeMs = decodeMs,
+ scoreMs = scoreMs,
+ elapsedMs = elapsedMs(startNs))
+ FetchResult(rows.iterator, keyMismatches.toSeq, metrics)
+ }
+
+ private def parquetRecordReader(
+ fileSchema: MessageType,
+ requestedParquetSchema: MessageType,
+ pageStore: PageReadStore): RecordReader[InternalRow] = {
+ val materializer = new InternalRowMaterializer(requestedParquetSchema, outputSchema, embeddingCol, distanceOrdinal)
+ new ColumnIOFactory().getColumnIO(requestedParquetSchema, fileSchema)
+ .getRecordReader(pageStore, materializer)
+ }
+
+ private def requestedSchema(fileSchema: MessageType): MessageType = {
+ val parquetFields =
+ if (vectorOnlyProjection) {
+ mutable.ArrayBuffer.empty[org.apache.parquet.schema.Type]
+ } else {
+ outputSchema.fields
+ .filterNot(field => isSyntheticOutputField(field.name))
+ .map(field => fieldByName(fileSchema, field.name))
+ .toBuffer[org.apache.parquet.schema.Type]
+ }
+ if ((verifyKeys || vectorOnlyProjection) && !parquetFields.exists(_.getName == HoodieRecord.RECORD_KEY_METADATA_FIELD)) {
+ parquetFields += fieldByName(fileSchema, HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ }
+ if (!parquetFields.exists(_.getName == embeddingCol)) {
+ parquetFields += fieldByName(fileSchema, embeddingCol)
+ }
+ new MessageType(fileSchema.getName, parquetFields.asJava)
+ }
+
+ private def isSyntheticOutputField(name: String): Boolean =
+ name == HoodieVectorSearchPlanBuilder.DISTANCE_COL ||
+ name == HoodieRecord.RECORD_KEY_METADATA_FIELD ||
+ name == HoodieRecord.PARTITION_PATH_METADATA_FIELD ||
+ name == HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL
+
+ private def fieldByName(schema: MessageType, name: String): org.apache.parquet.schema.Type =
+ schema.getFields.asScala.find(_.getName == name)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet schema is missing requested field '$name'."))
+
+ private def rowPosition(candidate: ScoredVectorPostingMatch): Long = {
+ val location = candidate.getLocation
+ if (location == null) {
+ throw new IllegalArgumentException(
+ s"Vector exact positional fetch requires a row-position locator for record ${candidate.getRecordKey}.")
+ }
+ location.getPosition
+ }
+
+ private def decodedRecordKey(row: InternalRow): String =
+ recordKeyOrdinal.flatMap(ordinal => Option(row.getUTF8String(ordinal))).map(_.toString).orNull
+
+ private def fillSyntheticFields(row: InternalRow, candidate: ScoredVectorPostingMatch): Unit = {
+ val liveLocation = Option(candidate.getLocation)
+ recordKeyOrdinal.foreach(row.update(_, UTF8String.fromString(candidate.getRecordKey)))
+ partitionPathOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getPartitionPath).orElse(Option(candidate.getPartitionPath)).getOrElse(""))))
+ fileGroupIdOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getFileId).orElse(Option(candidate.getFileGroupId)).getOrElse(""))))
+ }
+
+ private def fieldOrdinal(name: String): Option[Int] =
+ outputSchema.fields.indexWhere(_.name == name) match {
+ case -1 => None
+ case ordinal => Some(ordinal)
+ }
+
+ private def scoreVector(bytes: Array[Byte]): Double = scorer.scoreBytes(bytes)
+
+ private final class InternalRowMaterializer(
+ requestedParquetSchema: MessageType,
+ outputSchema: StructType,
+ embeddingCol: String,
+ distanceOrdinal: Int) extends RecordMaterializer[InternalRow] {
+
+ private val values = new Array[Any](outputSchema.length)
+ private var current: InternalRow = _
+ private var vectorBytes: Array[Byte] = _
+
+ private val root = new GroupConverter {
+ private val converters = requestedParquetSchema.getFields.asScala.map(converterFor).toIndexedSeq
+
+ override def getConverter(fieldIndex: Int): Converter =
+ converters(fieldIndex)
+
+ override def start(): Unit = {
+ var idx = 0
+ while (idx < values.length) {
+ values(idx) = null
+ idx += 1
+ }
+ vectorBytes = null
+ }
+
+ override def end(): Unit = {
+ if (vectorBytes == null) {
+ throw new IllegalStateException(s"Vector column '$embeddingCol' was not decoded for positional exact fetch.")
+ }
+ val scoreStartNs = System.nanoTime()
+ values(distanceOrdinal) = scoreVector(vectorBytes)
+ scoreMs += elapsedMs(scoreStartNs)
+ current = new GenericInternalRow(values.clone())
+ }
+ }
+
+ override def getCurrentRecord: InternalRow = current
+
+ override def getRootConverter: GroupConverter = root
+
+ private def converterFor(field: org.apache.parquet.schema.Type): PrimitiveConverter = {
+ val name = field.getName
+ new PrimitiveConverter {
+ override def addBinary(value: Binary): Unit = {
+ if (name == embeddingCol) {
+ vectorBytes = value.getBytes
+ } else {
+ setOutputValue(name, value.getBytes)
+ }
+ }
+
+ override def addInt(value: Int): Unit = setOutputValue(name, value)
+
+ override def addLong(value: Long): Unit = setOutputValue(name, value)
+
+ override def addFloat(value: Float): Unit = setOutputValue(name, value)
+
+ override def addDouble(value: Double): Unit = setOutputValue(name, value)
+
+ override def addBoolean(value: Boolean): Unit = setOutputValue(name, value)
+ }
+ }
+
+ private def setOutputValue(name: String, raw: Any): Unit = {
+ if (name != embeddingCol && name != HoodieVectorSearchPlanBuilder.DISTANCE_COL) {
+ val ordinal = outputSchema.fieldIndex(name)
+ values(ordinal) = outputSchema.fields(ordinal).dataType match {
+ case StringType => raw match {
+ case bytes: Array[Byte] => UTF8String.fromBytes(bytes)
+ case other => UTF8String.fromString(String.valueOf(other))
+ }
+ case BinaryType => raw.asInstanceOf[Array[Byte]]
+ case IntegerType => raw.asInstanceOf[Int]
+ case LongType => raw.asInstanceOf[Long]
+ case FloatType => raw.asInstanceOf[Float]
+ case DoubleType => raw.asInstanceOf[Double]
+ case BooleanType => raw.asInstanceOf[Boolean]
+ case ByteType => raw.asInstanceOf[Int].toByte
+ case ShortType => raw.asInstanceOf[Int].toShort
+ case other => throw new UnsupportedOperationException(
Review Comment:
🤖 non-blocking: the exact materializer only handles string/binary/int/long/float/double/boolean/byte/short leaf columns (and `converterFor` always returns a `PrimitiveConverter`), so any corpus column of timestamp, date, decimal, struct, array or map type makes exact mode fail inside the executor while brute force handles the same table; could this be validated at plan time with a clear message, or widened (the same list applies to `LogResidentVectorFetcher.toCatalyst`)?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-common/src/main/java/org/apache/hudi/common/index/vector/RaBitQEncoder.java:
##########
@@ -0,0 +1,588 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.common.index.vector;
+
+import java.io.Serializable;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Random;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * RaBitQ (Randomized Binary Quantization) encoder.
+ *
+ * <p>Encodes float vectors into 1-bit-per-dimension binary codes using a random
+ * orthogonal rotation matrix R. The rotation is deterministic given (seed, dimension),
+ * so R is stored as a seed — it never needs retraining.
+ *
+ * <p><b>Encoding (write path):</b>
+ * <ol>
+ * <li>Normalize: v̂ = v / ||v||</li>
+ * <li>Rotate: v_rot = R @ v̂</li>
+ * <li>Binarize: code = pack(sign(v_rot)) — D bits → ceil(D/8) bytes</li>
+ * <li>Scalar: s = ||v|| (stored alongside code; 1.0 if assume_normalized)</li>
+ * </ol>
+ *
+ * <p><b>Query scan (read path):</b>
+ * <ol>
+ * <li>Normalize + rotate the query vector (same R)</li>
+ * <li>Binarize to get q_bin</li>
+ * <li>For each database code b: Hamming(q_bin, b) → estimated cosine</li>
+ * <li>Re-rank top-R candidates with exact distance</li>
+ * </ol>
+ *
+ * <p>Memory: the rotation matrix is D×D floats = 4·D² bytes (~2.3 MB at D=768).
+ * It is built lazily on first use and reused across encode calls.
+ *
+ * <p>Thread-safe after the first call to {@link #encode} or {@link #encodeQuery}
+ * (the lazy init is synchronized).
+ */
+public final class RaBitQEncoder implements Serializable {
+ private static final long serialVersionUID = 1L;
+ private static final Map<RotationKey, float[][]> ROTATION_MATRIX_CACHE = new ConcurrentHashMap<>();
+ /** Neutral-factor thresholds for the current posting-block format (RFC-109 §3). */
+ private static final RaBitQFactorConfig FACTOR_CONFIG = RaBitQFactorConfig.defaults();
+ private final int dimension;
+ private final int bits;
+ private final long seed;
+ private final boolean assumeNormalized;
+
+ /** Rotation matrix, row-major. Populated lazily. */
+ private transient volatile float[][] rotMat;
+ public RaBitQEncoder(int dimension, long seed, boolean assumeNormalized) {
+ this(dimension, 1, seed, assumeNormalized);
+ }
+
+ public RaBitQEncoder(int dimension, int bits, long seed, boolean assumeNormalized) {
+ if (dimension <= 0) {
+ throw new IllegalArgumentException("Dimension must be positive, got: " + dimension);
+ }
+ if (bits <= 0 || bits > 8) {
+ throw new IllegalArgumentException("RaBitQ bits must be in [1, 8], got: " + bits);
+ }
+ this.dimension = dimension;
+ this.bits = bits;
+ this.seed = seed;
+ this.assumeNormalized = assumeNormalized;
+ }
+
+ /** Convenience constructor using default seed. */
+ public RaBitQEncoder(int dimension) {
+ this(dimension, 1, 42L, false);
+ }
+
+ // ---- encoding ----------------------------------------------------------
+
+ public QuantizedVector encode(float[] vector) {
+ validateVector(vector, "vector");
+ float norm = norm(vector);
+ float scalar = assumeNormalized ? 1.0f : norm;
+ float[] normalized = (norm == 0f || assumeNormalized)
+ ? vector
+ : normalize(vector, norm);
+ float[] rotated = rotate(normalized);
+ return new QuantizedVector(binarize(rotated), scalar);
+ }
+
+ /**
+ * Metric-neutral residual encoding used by MDT postings.
+ *
+ * <p>The optional {@code center} is the IVF centroid in the <b>original</b> vector space.
+ * When null, the vector is encoded relative to the origin.
+ */
+ public QuantizedVector encodeResidual(float[] vector, float[] center) {
Review Comment:
🤖 nit: the three `new QuantizedVector(...)` calls here pass 10 positional floats (with `centerRip` and `fRescale1` repeated in different slots), which is easy to misorder; would it make sense to add a `QuantizedVector.fromFactors(binaryCode, extendedCode, factors, bits)` factory so the slot mapping lives in one place?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/PositionalParquetFetcher.scala:
##########
@@ -0,0 +1,587 @@
+/*
+ * 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.spark.sql.hudi.analysis
+
+import org.apache.hudi.common.index.vector.{ScoredVectorPostingMatch, VectorIndexMdtSearchUtils, VectorPostingMatch}
+import org.apache.hudi.common.model.HoodieRecord
+import org.apache.hudi.common.schema.HoodieSchema
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileSystem, Path}
+import org.apache.parquet.column.page.PageReadStore
+import org.apache.parquet.hadoop.ParquetFileReader
+import org.apache.parquet.hadoop.metadata.{BlockMetaData, ColumnChunkMetaData, ParquetMetadata}
+import org.apache.parquet.internal.column.columnindex.OffsetIndex
+import org.apache.parquet.internal.filter2.columnindex.RowRanges
+import org.apache.parquet.io.{ColumnIOFactory, RecordReader}
+import org.apache.parquet.io.api.{Binary, Converter, GroupConverter, PrimitiveConverter, RecordMaterializer}
+import org.apache.parquet.schema.MessageType
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction.DistanceMetric
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+import org.slf4j.LoggerFactory
+
+import java.util.PrimitiveIterator
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+/**
+ * Fetches exact vector-search rerank candidates from Parquet by physical row position.
+ *
+ * The fetcher uses Parquet page indexes to turn candidate row positions into row ranges, so
+ * the reader can skip non-candidate pages instead of scanning whole Hudi file slices.
+ */
+private[analysis] final class PositionalParquetFetcher(
+ conf: Configuration,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Double],
+ metric: DistanceMetric.Value,
+ outputSchema: StructType,
+ failOnMissingOffsetIndex: Boolean,
+ vectorOnlyProjection: Boolean,
+ verifyKeys: Boolean) extends Serializable {
+
+ import PositionalParquetFetcher._
+
+ private val distanceOrdinal = outputSchema.fieldIndex(HoodieVectorSearchPlanBuilder.DISTANCE_COL)
+ private val recordKeyOrdinal = fieldOrdinal(HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ private val partitionPathOrdinal = fieldOrdinal(HoodieRecord.PARTITION_PATH_METADATA_FIELD)
+ private val fileGroupIdOrdinal = fieldOrdinal(HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL)
+ private val scorer = new VectorExactScorer(vectorSchema, queryVector, metric)
+ private var scoreMs: Long = 0L
+
+ def fetch(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch]): FetchResult = {
+ val startNs = System.nanoTime()
+ if (candidates.isEmpty) {
+ FetchResult(Iterator.empty, Seq.empty, FetchMetrics(elapsedMs = elapsedMs(startNs)))
+ } else {
+ fetchNonEmpty(filePath, candidates, startNs)
+ }
+ }
+
+ private def fetchNonEmpty(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch],
+ startNs: Long): FetchResult = {
+ val path = new Path(filePath)
+ configureRandomRead(conf)
+ val fs = path.getFileSystem(conf)
+ val fileLen = fs.getFileStatus(path).getLen
+ val fileCacheKey = s"$filePath#$fileLen"
+ val stats = Option(FileSystem.getStatistics(path.toUri.getScheme, fs.getClass))
+ val physicalBefore = stats.map(_.getBytesRead).getOrElse(0L)
+
+ val footerStartNs = System.nanoTime()
+ val footerResult = footer(conf, path, fileCacheKey)
+ val footerMs = elapsedMs(footerStartNs)
+ val metadata = footerResult.metadata
+ val fileSchema = metadata.getFileMetaData.getSchema
+ val requestedParquetSchema = requestedSchema(fileSchema)
+
+ val rowGroups = metadata.getBlocks.asScala.toIndexedSeq
+ val rowGroupStarts = cumulativeRowGroupStarts(rowGroups)
+ val byRowGroup: Map[Int, Seq[ScoredVectorPostingMatch]] =
+ candidates.groupBy(candidate => rowGroupOrdinal(rowGroupStarts, rowGroups, rowPosition(candidate)))
+
+ val rows = mutable.ArrayBuffer.empty[InternalRow]
+ val keyMismatches = mutable.ArrayBuffer.empty[ScoredVectorPostingMatch]
+ var rowGroupsSelected = 0
+ var pagesInSelectedRowGroups = 0L
+ var pagesSelected = 0L
+ var offsetIndexHits = 0L
+ var offsetIndexMissing = 0L
+ var pageBytesFetched = 0L
+ var rangedGets = 0L
+ var rowsDecoded = 0L
+ var rowsMaterialized = 0L
+ var offsetIndexMs = 0L
+ var decodeMs = 0L
+ scoreMs = 0L
+
+ val reader = new ParquetFileReader(conf, path, metadata)
+ try {
+ reader.setRequestedSchema(requestedParquetSchema)
+ byRowGroup.toSeq.sortBy(_._1).foreach { case (rgOrdinal, rgCandidates) =>
+ val block = rowGroups(rgOrdinal)
+ val rowGroupStart = rowGroupStarts(rgOrdinal)
+ val embeddingChunk = findColumn(block, embeddingCol)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet file '$filePath' is missing vector column '$embeddingCol'."))
+
+ val offsetStartNs = System.nanoTime()
+ val embeddingOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, embeddingChunk)
+ offsetIndexMs += elapsedMs(offsetStartNs)
+ val relativeRows = rgCandidates.map(candidate => rowPosition(candidate) - rowGroupStart).sorted
+ val candidateRows = relativeRows.toSet
+ val candidateByRelativeRow = rgCandidates
+ .map(candidate => rowPosition(candidate) - rowGroupStart -> candidate)
+ .toMap
+
+ if (embeddingOffsetIndex == null && failOnMissingOffsetIndex) {
+ offsetIndexMissing += 1
+ throw new IllegalStateException(
+ s"Parquet file '$filePath' has no OffsetIndex for vector column '$embeddingCol'; positional exact fetch cannot honor pages.")
+ }
+
+ rowGroupsSelected += 1
+ val (pageStore, rowRanges) = if (embeddingOffsetIndex == null) {
+ offsetIndexMissing += 1
+ LOG.warn(
+ "Parquet file {} rowGroup={} has no OffsetIndex for vector column {}; falling back to row-group fetch.",
+ filePath,
+ Int.box(rgOrdinal),
+ embeddingCol)
+ pagesInSelectedRowGroups += 0L
+ pagesSelected += 0L
+ rangedGets += 1L
+ pageBytesFetched += requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .map(_.getTotalSize)
+ .sum
+ (reader.readRowGroup(rgOrdinal), RowRanges.createSingle(block.getRowCount))
+ } else {
+ offsetIndexHits += 1
+ pagesInSelectedRowGroups += embeddingOffsetIndex.getPageCount
+ val selectedVectorPages = selectedPages(embeddingOffsetIndex, block.getRowCount, relativeRows)
+ pagesSelected += selectedVectorPages.size
+ val selectedRanges = requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .flatMap { chunk =>
+ val oiStartNs = System.nanoTime()
+ val chunkOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, chunk)
+ offsetIndexMs += elapsedMs(oiStartNs)
+ if (chunkOffsetIndex == null) {
+ offsetIndexMissing += 1
+ Seq.empty
+ } else {
+ offsetIndexHits += 1
+ selectedPages(chunkOffsetIndex, block.getRowCount, relativeRows).map { pageOrdinal =>
+ val start = chunkOffsetIndex.getOffset(pageOrdinal)
+ PageRange(start, start + chunkOffsetIndex.getCompressedPageSize(pageOrdinal))
+ }
+ }
+ }
+ val mergedRanges = mergeRanges(selectedRanges.sortBy(_.start).toSeq)
+ rangedGets += mergedRanges.size
+ pageBytesFetched += mergedRanges.iterator.map(range => range.end - range.start).sum
+ val ranges = RowRanges.create(block.getRowCount, primitiveIntIterator(selectedVectorPages), embeddingOffsetIndex)
+ (reader.readFilteredRowGroup(rgOrdinal, ranges), ranges)
+ }
+
+ val decodeStartNs = System.nanoTime()
+ val recordReader = parquetRecordReader(fileSchema, requestedParquetSchema, pageStore)
+ val rowIterator = rowRanges.iterator()
+ while (rowIterator.hasNext) {
+ val relativeRow = rowIterator.nextLong()
+ val row = recordReader.read()
+ rowsDecoded += 1
+ candidateByRelativeRow.get(relativeRow).foreach { candidate =>
+ if (verifyKeys && decodedRecordKey(row) != candidate.getRecordKey) {
+ keyMismatches += candidate
+ } else {
+ fillSyntheticFields(row, candidate)
+ rows += row
+ rowsMaterialized += 1
+ }
+ }
+ }
+ decodeMs += elapsedMs(decodeStartNs)
+ }
+ } finally {
+ reader.close()
+ }
+
+ val physicalAfter = stats.map(_.getBytesRead).getOrElse(physicalBefore)
+ val metrics = FetchMetrics(
+ rerankCandidates = candidates.size,
+ rowGroupsTotal = rowGroups.size,
+ rowGroupsSelected = rowGroupsSelected,
+ pagesInSelectedRowGroups = pagesInSelectedRowGroups,
+ pagesSelected = pagesSelected,
+ offsetIndexHits = offsetIndexHits,
+ offsetIndexMissing = offsetIndexMissing,
+ rangedGets = rangedGets,
+ pageBytesFetched = pageBytesFetched,
+ rowsDecoded = rowsDecoded,
+ rowsMaterialized = rowsMaterialized,
+ fetchKeyMismatch = keyMismatches.size,
+ physicalBytesRead = math.max(0L, physicalAfter - physicalBefore),
+ footerReads = if (footerResult.cacheHit) 0 else 1,
+ footerCacheHits = if (footerResult.cacheHit) 1 else 0,
+ footerMs = footerMs,
+ offsetIndexMs = offsetIndexMs,
+ fetchWaitMs = 0L,
+ decodeMs = decodeMs,
+ scoreMs = scoreMs,
+ elapsedMs = elapsedMs(startNs))
+ FetchResult(rows.iterator, keyMismatches.toSeq, metrics)
+ }
+
+ private def parquetRecordReader(
+ fileSchema: MessageType,
+ requestedParquetSchema: MessageType,
+ pageStore: PageReadStore): RecordReader[InternalRow] = {
+ val materializer = new InternalRowMaterializer(requestedParquetSchema, outputSchema, embeddingCol, distanceOrdinal)
+ new ColumnIOFactory().getColumnIO(requestedParquetSchema, fileSchema)
+ .getRecordReader(pageStore, materializer)
+ }
+
+ private def requestedSchema(fileSchema: MessageType): MessageType = {
+ val parquetFields =
+ if (vectorOnlyProjection) {
+ mutable.ArrayBuffer.empty[org.apache.parquet.schema.Type]
+ } else {
+ outputSchema.fields
+ .filterNot(field => isSyntheticOutputField(field.name))
+ .map(field => fieldByName(fileSchema, field.name))
+ .toBuffer[org.apache.parquet.schema.Type]
+ }
+ if ((verifyKeys || vectorOnlyProjection) && !parquetFields.exists(_.getName == HoodieRecord.RECORD_KEY_METADATA_FIELD)) {
+ parquetFields += fieldByName(fileSchema, HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ }
+ if (!parquetFields.exists(_.getName == embeddingCol)) {
+ parquetFields += fieldByName(fileSchema, embeddingCol)
+ }
+ new MessageType(fileSchema.getName, parquetFields.asJava)
+ }
+
+ private def isSyntheticOutputField(name: String): Boolean =
+ name == HoodieVectorSearchPlanBuilder.DISTANCE_COL ||
+ name == HoodieRecord.RECORD_KEY_METADATA_FIELD ||
+ name == HoodieRecord.PARTITION_PATH_METADATA_FIELD ||
+ name == HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL
+
+ private def fieldByName(schema: MessageType, name: String): org.apache.parquet.schema.Type =
+ schema.getFields.asScala.find(_.getName == name)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet schema is missing requested field '$name'."))
+
+ private def rowPosition(candidate: ScoredVectorPostingMatch): Long = {
+ val location = candidate.getLocation
+ if (location == null) {
+ throw new IllegalArgumentException(
+ s"Vector exact positional fetch requires a row-position locator for record ${candidate.getRecordKey}.")
+ }
+ location.getPosition
+ }
+
+ private def decodedRecordKey(row: InternalRow): String =
+ recordKeyOrdinal.flatMap(ordinal => Option(row.getUTF8String(ordinal))).map(_.toString).orNull
+
+ private def fillSyntheticFields(row: InternalRow, candidate: ScoredVectorPostingMatch): Unit = {
+ val liveLocation = Option(candidate.getLocation)
+ recordKeyOrdinal.foreach(row.update(_, UTF8String.fromString(candidate.getRecordKey)))
+ partitionPathOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getPartitionPath).orElse(Option(candidate.getPartitionPath)).getOrElse(""))))
+ fileGroupIdOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getFileId).orElse(Option(candidate.getFileGroupId)).getOrElse(""))))
+ }
+
+ private def fieldOrdinal(name: String): Option[Int] =
+ outputSchema.fields.indexWhere(_.name == name) match {
+ case -1 => None
+ case ordinal => Some(ordinal)
+ }
+
+ private def scoreVector(bytes: Array[Byte]): Double = scorer.scoreBytes(bytes)
+
+ private final class InternalRowMaterializer(
+ requestedParquetSchema: MessageType,
+ outputSchema: StructType,
+ embeddingCol: String,
+ distanceOrdinal: Int) extends RecordMaterializer[InternalRow] {
+
+ private val values = new Array[Any](outputSchema.length)
+ private var current: InternalRow = _
+ private var vectorBytes: Array[Byte] = _
+
+ private val root = new GroupConverter {
+ private val converters = requestedParquetSchema.getFields.asScala.map(converterFor).toIndexedSeq
+
+ override def getConverter(fieldIndex: Int): Converter =
+ converters(fieldIndex)
+
+ override def start(): Unit = {
+ var idx = 0
+ while (idx < values.length) {
+ values(idx) = null
+ idx += 1
+ }
+ vectorBytes = null
+ }
+
+ override def end(): Unit = {
+ if (vectorBytes == null) {
+ throw new IllegalStateException(s"Vector column '$embeddingCol' was not decoded for positional exact fetch.")
Review Comment:
🤖 non-blocking: a candidate whose live row now has a NULL embedding (indexed, later upserted to NULL, queried under `vector.freshness.policy=warn`) is key-located and lands here, and the throw fails the whole query whereas brute force just filters nulls; would skipping the row (and counting it) be the better outcome, here and in `LogResidentVectorFetcher.embeddingBytes`?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
##########
hudi-common/src/main/java/org/apache/hudi/common/index/vector/VectorIndexMdtSearchUtils.java:
##########
@@ -0,0 +1,1310 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.hudi.common.index.vector;
+
+import org.apache.hudi.avro.model.HoodieVectorIndexPostingBlock;
+import org.apache.hudi.common.data.HoodieData;
+import org.apache.hudi.common.data.HoodieListData;
+import org.apache.hudi.common.data.HoodiePairData;
+import org.apache.hudi.common.model.HoodieRecord;
+import org.apache.hudi.common.model.HoodieRecordGlobalLocation;
+import org.apache.hudi.common.util.Option;
+import org.apache.hudi.common.util.collection.Pair;
+import org.apache.hudi.metadata.HoodieMetadataPayload;
+import org.apache.hudi.metadata.HoodieTableMetadata;
+import org.apache.hudi.metadata.RawKey;
+import org.apache.hudi.metadata.VectorIndexMetadataKey;
+import org.apache.hudi.metadata.VectorPostingPrefixRawKey;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.lang.management.GarbageCollectorMXBean;
+import java.lang.management.ManagementFactory;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+
+/**
+ * Shared helper for MDT-native vector posting lookup and approximate candidate reduction.
+ *
+ * <p>This provides the missing second-stage building block after coarse IVF cluster pruning:
+ * exact HFile lookups of cluster metadata, prefix scans over {@code P|gen|cluster|shard|},
+ * approximate RaBitQ scoring, and direct file-slice targeting from posting payload metadata.
+ */
+public final class VectorIndexMdtSearchUtils {
+
+ private static final Logger LOG = LoggerFactory.getLogger(VectorIndexMdtSearchUtils.class);
+ private static final int TOP_K_REDUCER_KEY = 0;
+ private static final Comparator<ScoredVectorPostingMatch> BEST_FIRST = Comparator
+ .comparingDouble(ScoredVectorPostingMatch::getApproxDistance)
+ .thenComparing(ScoredVectorPostingMatch::getRecordKey);
+
+ private VectorIndexMdtSearchUtils() {
+ }
+
+ public static Map<Integer, Integer> readClusterShardCounts(HoodieTableMetadata metadataTable,
+ String indexPartition,
+ int generationId,
+ Collection<Integer> clusterIds) {
+ if (clusterIds == null || clusterIds.isEmpty()) {
+ return Collections.emptyMap();
+ }
+
+ Map<Integer, Integer> shardCounts = new HashMap<>();
+ for (Integer clusterId : clusterIds) {
+ shardCounts.put(clusterId, 1);
+ }
+ return shardCounts;
+ }
+
+ public static Map<Integer, Set<String>> readClusterFileGroups(HoodieTableMetadata metadataTable,
+ String indexPartition,
+ int generationId,
+ Collection<Integer> clusterIds,
+ Collection<String> partitionPaths,
+ boolean shouldLoadInMemory) {
+ if (clusterIds == null || clusterIds.isEmpty()) {
+ return Collections.emptyMap();
+ }
+
+ return Collections.emptyMap();
+ }
+
+ public static List<VectorPostingPrefixRawKey> buildPostingPrefixes(int generationId,
+ Map<Integer, Integer> clusterShardCounts) {
+ if (clusterShardCounts == null || clusterShardCounts.isEmpty()) {
+ return Collections.emptyList();
+ }
+
+ List<Map.Entry<Integer, Integer>> clusters = new ArrayList<>(clusterShardCounts.entrySet());
+ clusters.sort(Map.Entry.comparingByKey());
+
+ List<VectorPostingPrefixRawKey> prefixes = new ArrayList<>();
+ for (Map.Entry<Integer, Integer> entry : clusters) {
+ int clusterId = entry.getKey();
+ int shardCount = Math.max(1, entry.getValue());
+ for (int shardId = 0; shardId < shardCount; shardId++) {
+ prefixes.add(new VectorPostingPrefixRawKey(generationId, clusterId, shardId));
+ }
+ }
+ return prefixes;
+ }
+
+ public static HoodieData<VectorPostingMatch> readPostingMatches(HoodieTableMetadata metadataTable,
+ String indexPartition,
+ int generationId,
+ Map<Integer, Integer> clusterShardCounts,
+ boolean shouldLoadInMemory) {
+ List<VectorPostingPrefixRawKey> postingPrefixes = buildPostingPrefixes(generationId, clusterShardCounts);
+ if (postingPrefixes.isEmpty()) {
+ return HoodieListData.eager(Collections.emptyList());
+ }
+
+ List<RawKey> rawKeys = new ArrayList<>(postingPrefixes);
+ HoodieData<VectorPostingMatch> matches = metadataTable.getRecordsByKeyPrefixes(HoodieListData.eager(rawKeys), indexPartition, shouldLoadInMemory)
+ .flatMap(record -> {
+ Option<Object> infoOpt = getVectorInfo(record);
+ if (!infoOpt.isPresent()) {
+ return Collections.<VectorPostingMatch>emptyIterator();
+ }
+
+ if (record.getData().isDeleted()) {
+ String deletedRecordKey = VectorIndexMetadataKey.postingRecordKey(record.getRecordKey());
+ return deletedRecordKey == null
+ ? Collections.<VectorPostingMatch>emptyIterator()
+ : Collections.singletonList(VectorPostingMatch.tombstone(deletedRecordKey, VectorPostingSerde.parsePostingKey(record.getRecordKey()))).iterator();
+ }
+
+ Object info = infoOpt.get();
+ int[] keyComponents = VectorPostingSerde.parsePostingKey(record.getRecordKey());
+ if (VectorPostingSerde.isPostingTombstoneInfo(info)) {
+ String deletedRecordKey = VectorIndexMetadataKey.postingRecordKey(record.getRecordKey());
+ return deletedRecordKey == null
+ ? Collections.<VectorPostingMatch>emptyIterator()
+ : Collections.singletonList(VectorPostingMatch.tombstone(deletedRecordKey, keyComponents)).iterator();
+ }
+ if (VectorPostingSerde.isPostingBlockInfo(info)) {
+ PostingBlockView view = new PostingBlockView(asPostingBlock(info));
+ List<VectorPostingMatch> blockMatches = new ArrayList<>(view.numVectors());
+ for (int vectorIndex = 0; vectorIndex < view.numVectors(); vectorIndex++) {
+ PostingBlockView.RowLocator rowLocator = view.rowLocator(vectorIndex);
+ blockMatches.add(new VectorPostingMatch(
Review Comment:
🤖 nit: this 13-arg `VectorPostingMatch` construction from a `PostingBlockView` row is duplicated verbatim in `scanBlockPartition`; could you pull it into a `VectorPostingMatch.fromBlockRow(view, vectorIndex, clusterId, shardId)` helper so the scalar-factor slot mapping is defined once?
<sub><i>⚠️ AI-generated; verify before applying. React 👍/👎 to flag quality.</i></sub>
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on PR #20107:
URL: https://github.com/apache/hudi/pull/20107#issuecomment-5924132829
> is this the latest PR.. like the new one for us to review @chrevanthreddy ..
>
> @vinishjail97 is going to work with you closely and get this over.. test this within Onehouse as well parallely..
>
> @yihua and I can help reviews as needed
Yes. This is the refactor to separate the core changes, vector changes and hfile cache related changes for easier review.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(index): refactor RFC-109 vector bootstrap and Spark search [hudi]
Posted by "chrevanthreddy (via GitHub)" <gi...@apache.org>.
chrevanthreddy commented on code in PR #20107:
URL: https://github.com/apache/hudi/pull/20107#discussion_r4384293478
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorExactSearchPlanner.scala:
##########
@@ -0,0 +1,390 @@
+/*
+ * 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.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.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, 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}
+
+/** Exact reranking and file-slice resolution for the IVF/RaBitQ search plan. */
+private[analysis] object VectorExactSearchPlanner extends SparkAdapterSupport {
+ import HoodieVectorSearchPlanBuilder._
+ import IvfRaBitQMdtSearchAlgorithm.{isLogResidentCandidate, isPartitionedRecordIndex, toVectorMetric, VectorSearchPlanContext}
+
+ private val LOG = LoggerFactory.getLogger(getClass)
+
+ def buildRaBitQExactCandidateDf(
+ spark: SparkSession,
+ ctx: VectorSearchPlanContext,
+ outputSchema: StructType,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Float],
+ exactQueryVector: Array[Double],
+ k: Int): DataFrame = {
+ val planStartNs = System.nanoTime()
+ val storageConf = HadoopFSUtils.getStorageConf(spark.sessionState.newHadoopConf())
+ val cacheLoadStartNs = 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 topClusters = cache.findTopClusters(queryVector, numProbes, metric)
+ val topClusterSeq = topClusters.toSeq
+ val cacheLoadMs = elapsedMs(cacheLoadStartNs)
+
+ // Shard counts from cache — no extra MDT IO.
+ val shardCounts = cache.getShardCounts(topClusters)
+ val shardCountsScala = shardCounts.asScala.map { case (clusterId, count) =>
+ clusterId.intValue() -> count.intValue()
+ }.toMap
+ val postingPrefixCount = VectorIndexMdtSearchUtils.buildPostingPrefixes(cache.getGenerationId, shardCounts).size()
+ LOG.info(
+ s"[vector_search][plan] basePath=${ctx.basePath} indexPartition=${ctx.indexPartition} cacheLoadMs=$cacheLoadMs " +
+ s"generationId=${cache.getGenerationId} dimension=${cache.getDimension} numClusters=${cache.numClusters()} " +
+ s"numProbes=$numProbes refineFactor=${VectorIndexOptions.getRefineFactor(mergedOptions.asJava)} residualEncoding=$residualEncoding " +
+ s"topClusters=${summarizeInts(topClusterSeq)} shardCounts=${summarizeShardCounts(shardCountsScala)} " +
+ s"postingPrefixes=$postingPrefixCount")
+ if (shardCounts.isEmpty) {
+ sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+ } else {
+ val refineK = math.max(k, k * VectorIndexOptions.getRefineFactor(mergedOptions.asJava))
+ val emptyDf = sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+
+ 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, refineK)
+ val topCandidateMaterializeStartNs = System.nanoTime()
+ val materializedCandidates = try {
+ topCandidates.collectAsList().asScala.toSeq
+ } finally {
+ topCandidates.unpersistWithDependencies()
+ }
+ val topCandidateList = materializedCandidates
+ .filter(candidate => candidate.getFileGroupId != null && candidate.getFileGroupId.nonEmpty)
+ val candidatesWithoutLocator = materializedCandidates.size - topCandidateList.size
+ LOG.info(
+ s"[vector_search][stage][materialize_top_candidates] refineK=$refineK rerankCandidates=${materializedCandidates.size} " +
+ s"candidatesWithLocators=${topCandidateList.size} candidatesWithoutLocators=$candidatesWithoutLocator " +
+ s"distinctFileGroups=${topCandidateList.map(_.getFileGroupId).distinct.size} elapsedMs=${elapsedMs(topCandidateMaterializeStartNs)}")
+ if (candidatesWithoutLocator > 0) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search found $candidatesWithoutLocator rerank candidates without file-group locators. " +
+ s"Exact rerank requires per-candidate posting locators.")
+ }
+
+ // RFC-109 RLI finalist arbiter. Exact mode uses refreshed RLI locations for every live
+ // candidate: SERVE retains positional trust; STALE is resolved by key in the live base file.
+ val arbiterEnabled = VectorIndexOptions.isFinalistArbiterEnabled(mergedOptions.asJava)
+ val serveCandidates: Seq[ScoredVectorPostingMatch] =
+ if (arbiterEnabled) {
+ val arb = VectorIndexMdtSearchUtils.arbitrateMaterializedFinalists(
+ ctx.metadataTable, topCandidateList.asJava, isPartitionedRecordIndex(ctx))
+ val staleFailPolicy = VectorIndexOptions.isStaleLocatorPolicyFail(mergedOptions.asJava)
+ LOG.info(
+ s"[vector_search][exact][arbiter] arbiterExclusions{stale=${arb.staleCount}, deleted=${arb.deletedCount}} " +
+ s"serve=${arb.serve.size} stalePolicy=${if (staleFailPolicy) "fail" else "fallback"}")
+ if (staleFailPolicy && !arb.stale.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search rejected ${arb.staleCount} stale finalist(s) because " +
+ "vector.stale.locator.policy=FAIL.")
+ }
+ arb.serve.asScala.toSeq ++ arb.stale.asScala.toSeq
+ } else {
+ topCandidateList
+ }
+
+ if (serveCandidates.isEmpty) {
+ emptyDf
+ } else {
+ val candidatesByPartition = serveCandidates
+ .groupBy(candidate => Option(candidate.getLocation).map(_.getPartitionPath).getOrElse(""))
+ .map { case (partitionPath, candidates) =>
+ partitionPath -> candidates.flatMap(candidate => Option(candidate.getLocation).map(_.getFileId)).toSet
+ }
+ val fileSliceMap = preResolveFileSlices(
+ ctx.metadataTable, ctx.metaClient, ctx.timeline, ctx.latestInstantTime, candidatesByPartition)
+ if (fileSliceMap.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search resolved zero file slices for ${serveCandidates.size} rerank candidates " +
+ s"across ${candidatesByPartition.values.map(_.size).sum} file groups. Candidate locators are present, so this indicates " +
+ s"a file-slice resolution bug or stale vector index.")
+ }
+ val bFileSliceMap = spark.sparkContext.broadcast(fileSliceMap)
+ // MOR log-resident candidates carry their raw embedding in a log block, not the base Parquet
+ // file, so the positional fetcher cannot read them. Split them out and materialize them via a
+ // merged (base + log) read; base-resident candidates keep the fast positional path.
+ val (logResidentCandidates, baseCandidates) = serveCandidates.partition { candidate =>
+ bFileSliceMap.value.get(candidate.getLocation.getFileId).exists { slice =>
+ isLogResidentCandidate(candidate, slice.getBaseInstantTime, slice.getLogFiles.findAny().isPresent)
+ }
+ }
+ val exactFetchParallelism = math.max(
+ 1,
+ math.min(refineK, math.max(1, spark.sparkContext.defaultParallelism * 4)))
+ val groupedCandidates = baseCandidates.groupBy(_.getLocation.getFileId).toSeq.sortBy(_._1)
+ val targetBatchesPerGroup = math.max(1, exactFetchParallelism / math.max(1, groupedCandidates.size))
+ val fetchBatches = groupedCandidates.flatMap { case (fgId, candidates) =>
+ val sorted = candidates.sortBy(candidate => candidate.getLocation.getPosition)
+ val batchSize = math.max(1, math.ceil(sorted.size.toDouble / targetBatchesPerGroup).toInt)
+ sorted.grouped(batchSize).map(batch => (fgId, batch)).toSeq
+ }
+ val byFetchBatch: RDD[(String, Seq[ScoredVectorPostingMatch])] =
+ if (fetchBatches.isEmpty) {
+ spark.sparkContext.emptyRDD[(String, Seq[ScoredVectorPostingMatch])]
+ } else {
+ spark.sparkContext.parallelize(fetchBatches, math.min(exactFetchParallelism, fetchBatches.size))
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read_plan] rerankCandidates=${serveCandidates.size} " +
+ s"candidateGroups=${groupedCandidates.size} resolvedSlices=${fileSliceMap.size} " +
+ s"fileGroupPartitions=${byFetchBatch.getNumPartitions} exactFetchParallelism=$exactFetchParallelism")
+
+ val vectorOnlyProjection = mergedOptions.get(VECTOR_ONLY_PROJECTION_OPT)
+ .orElse(spark.conf.getOption(VECTOR_ONLY_PROJECTION_SPARK_CONF))
+ .exists(_.equalsIgnoreCase("true"))
+ val fetchSkippedLogResidentAcc: LongAccumulator =
+ spark.sparkContext.longAccumulator("vector_fetch_skipped_log_resident")
+ val fetchKeyMismatchAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_fetch_key_mismatch")
+ val rowsMaterializedAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_rows_materialized")
+ val exactRows: RDD[InternalRow] = byFetchBatch.mapPartitions { batches =>
+ val startNs = System.nanoTime()
+ val rows = ArrayBuffer.empty[InternalRow]
+ var batchCount = 0
+ var resolvedSlices = 0
+ var missingSlices = 0
+ var candidateKeys = 0L
+ var baseFiles = 0
+ var logFiles = 0L
+ var totalSliceBytes = 0L
+ var staleCandidates = 0L
+ var fetchSkippedLogResident = 0L
+ var fetchKeyMismatch = 0L
+ var keyLookupRowsDecoded = 0L
+ var metrics = PositionalParquetFetcher.FetchMetrics()
+ val fetcher = new PositionalParquetFetcher(
+ storageConf.newInstance().unwrap(),
+ embeddingCol,
+ vectorSchema,
+ exactQueryVector,
+ ctx.metric,
Review Comment:
**Fixed** `3cde5283ea`. The closure captures only the distance metric.
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/VectorExactSearchPlanner.scala:
##########
@@ -0,0 +1,390 @@
+/*
+ * 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.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.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, 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}
+
+/** Exact reranking and file-slice resolution for the IVF/RaBitQ search plan. */
+private[analysis] object VectorExactSearchPlanner extends SparkAdapterSupport {
+ import HoodieVectorSearchPlanBuilder._
+ import IvfRaBitQMdtSearchAlgorithm.{isLogResidentCandidate, isPartitionedRecordIndex, toVectorMetric, VectorSearchPlanContext}
+
+ private val LOG = LoggerFactory.getLogger(getClass)
+
+ def buildRaBitQExactCandidateDf(
+ spark: SparkSession,
+ ctx: VectorSearchPlanContext,
+ outputSchema: StructType,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Float],
+ exactQueryVector: Array[Double],
+ k: Int): DataFrame = {
+ val planStartNs = System.nanoTime()
+ val storageConf = HadoopFSUtils.getStorageConf(spark.sessionState.newHadoopConf())
+ val cacheLoadStartNs = 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 topClusters = cache.findTopClusters(queryVector, numProbes, metric)
+ val topClusterSeq = topClusters.toSeq
+ val cacheLoadMs = elapsedMs(cacheLoadStartNs)
+
+ // Shard counts from cache — no extra MDT IO.
+ val shardCounts = cache.getShardCounts(topClusters)
+ val shardCountsScala = shardCounts.asScala.map { case (clusterId, count) =>
+ clusterId.intValue() -> count.intValue()
+ }.toMap
+ val postingPrefixCount = VectorIndexMdtSearchUtils.buildPostingPrefixes(cache.getGenerationId, shardCounts).size()
+ LOG.info(
+ s"[vector_search][plan] basePath=${ctx.basePath} indexPartition=${ctx.indexPartition} cacheLoadMs=$cacheLoadMs " +
+ s"generationId=${cache.getGenerationId} dimension=${cache.getDimension} numClusters=${cache.numClusters()} " +
+ s"numProbes=$numProbes refineFactor=${VectorIndexOptions.getRefineFactor(mergedOptions.asJava)} residualEncoding=$residualEncoding " +
+ s"topClusters=${summarizeInts(topClusterSeq)} shardCounts=${summarizeShardCounts(shardCountsScala)} " +
+ s"postingPrefixes=$postingPrefixCount")
+ if (shardCounts.isEmpty) {
+ sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+ } else {
+ val refineK = math.max(k, k * VectorIndexOptions.getRefineFactor(mergedOptions.asJava))
+ val emptyDf = sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(
+ spark, spark.sparkContext.emptyRDD[InternalRow], outputSchema)
+
+ 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, refineK)
+ val topCandidateMaterializeStartNs = System.nanoTime()
+ val materializedCandidates = try {
+ topCandidates.collectAsList().asScala.toSeq
+ } finally {
+ topCandidates.unpersistWithDependencies()
+ }
+ val topCandidateList = materializedCandidates
+ .filter(candidate => candidate.getFileGroupId != null && candidate.getFileGroupId.nonEmpty)
+ val candidatesWithoutLocator = materializedCandidates.size - topCandidateList.size
+ LOG.info(
+ s"[vector_search][stage][materialize_top_candidates] refineK=$refineK rerankCandidates=${materializedCandidates.size} " +
+ s"candidatesWithLocators=${topCandidateList.size} candidatesWithoutLocators=$candidatesWithoutLocator " +
+ s"distinctFileGroups=${topCandidateList.map(_.getFileGroupId).distinct.size} elapsedMs=${elapsedMs(topCandidateMaterializeStartNs)}")
+ if (candidatesWithoutLocator > 0) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search found $candidatesWithoutLocator rerank candidates without file-group locators. " +
+ s"Exact rerank requires per-candidate posting locators.")
+ }
+
+ // RFC-109 RLI finalist arbiter. Exact mode uses refreshed RLI locations for every live
+ // candidate: SERVE retains positional trust; STALE is resolved by key in the live base file.
+ val arbiterEnabled = VectorIndexOptions.isFinalistArbiterEnabled(mergedOptions.asJava)
+ val serveCandidates: Seq[ScoredVectorPostingMatch] =
+ if (arbiterEnabled) {
+ val arb = VectorIndexMdtSearchUtils.arbitrateMaterializedFinalists(
+ ctx.metadataTable, topCandidateList.asJava, isPartitionedRecordIndex(ctx))
+ val staleFailPolicy = VectorIndexOptions.isStaleLocatorPolicyFail(mergedOptions.asJava)
+ LOG.info(
+ s"[vector_search][exact][arbiter] arbiterExclusions{stale=${arb.staleCount}, deleted=${arb.deletedCount}} " +
+ s"serve=${arb.serve.size} stalePolicy=${if (staleFailPolicy) "fail" else "fallback"}")
+ if (staleFailPolicy && !arb.stale.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search rejected ${arb.staleCount} stale finalist(s) because " +
+ "vector.stale.locator.policy=FAIL.")
+ }
+ arb.serve.asScala.toSeq ++ arb.stale.asScala.toSeq
+ } else {
+ topCandidateList
+ }
+
+ if (serveCandidates.isEmpty) {
+ emptyDf
+ } else {
+ val candidatesByPartition = serveCandidates
+ .groupBy(candidate => Option(candidate.getLocation).map(_.getPartitionPath).getOrElse(""))
+ .map { case (partitionPath, candidates) =>
+ partitionPath -> candidates.flatMap(candidate => Option(candidate.getLocation).map(_.getFileId)).toSet
+ }
+ val fileSliceMap = preResolveFileSlices(
+ ctx.metadataTable, ctx.metaClient, ctx.timeline, ctx.latestInstantTime, candidatesByPartition)
+ if (fileSliceMap.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact search resolved zero file slices for ${serveCandidates.size} rerank candidates " +
+ s"across ${candidatesByPartition.values.map(_.size).sum} file groups. Candidate locators are present, so this indicates " +
+ s"a file-slice resolution bug or stale vector index.")
+ }
+ val bFileSliceMap = spark.sparkContext.broadcast(fileSliceMap)
+ // MOR log-resident candidates carry their raw embedding in a log block, not the base Parquet
+ // file, so the positional fetcher cannot read them. Split them out and materialize them via a
+ // merged (base + log) read; base-resident candidates keep the fast positional path.
+ val (logResidentCandidates, baseCandidates) = serveCandidates.partition { candidate =>
+ bFileSliceMap.value.get(candidate.getLocation.getFileId).exists { slice =>
+ isLogResidentCandidate(candidate, slice.getBaseInstantTime, slice.getLogFiles.findAny().isPresent)
+ }
+ }
+ val exactFetchParallelism = math.max(
+ 1,
+ math.min(refineK, math.max(1, spark.sparkContext.defaultParallelism * 4)))
+ val groupedCandidates = baseCandidates.groupBy(_.getLocation.getFileId).toSeq.sortBy(_._1)
+ val targetBatchesPerGroup = math.max(1, exactFetchParallelism / math.max(1, groupedCandidates.size))
+ val fetchBatches = groupedCandidates.flatMap { case (fgId, candidates) =>
+ val sorted = candidates.sortBy(candidate => candidate.getLocation.getPosition)
+ val batchSize = math.max(1, math.ceil(sorted.size.toDouble / targetBatchesPerGroup).toInt)
+ sorted.grouped(batchSize).map(batch => (fgId, batch)).toSeq
+ }
+ val byFetchBatch: RDD[(String, Seq[ScoredVectorPostingMatch])] =
+ if (fetchBatches.isEmpty) {
+ spark.sparkContext.emptyRDD[(String, Seq[ScoredVectorPostingMatch])]
+ } else {
+ spark.sparkContext.parallelize(fetchBatches, math.min(exactFetchParallelism, fetchBatches.size))
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read_plan] rerankCandidates=${serveCandidates.size} " +
+ s"candidateGroups=${groupedCandidates.size} resolvedSlices=${fileSliceMap.size} " +
+ s"fileGroupPartitions=${byFetchBatch.getNumPartitions} exactFetchParallelism=$exactFetchParallelism")
+
+ val vectorOnlyProjection = mergedOptions.get(VECTOR_ONLY_PROJECTION_OPT)
+ .orElse(spark.conf.getOption(VECTOR_ONLY_PROJECTION_SPARK_CONF))
+ .exists(_.equalsIgnoreCase("true"))
+ val fetchSkippedLogResidentAcc: LongAccumulator =
+ spark.sparkContext.longAccumulator("vector_fetch_skipped_log_resident")
+ val fetchKeyMismatchAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_fetch_key_mismatch")
+ val rowsMaterializedAcc: LongAccumulator = spark.sparkContext.longAccumulator("vector_rows_materialized")
+ val exactRows: RDD[InternalRow] = byFetchBatch.mapPartitions { batches =>
+ val startNs = System.nanoTime()
+ val rows = ArrayBuffer.empty[InternalRow]
+ var batchCount = 0
+ var resolvedSlices = 0
+ var missingSlices = 0
+ var candidateKeys = 0L
+ var baseFiles = 0
+ var logFiles = 0L
+ var totalSliceBytes = 0L
+ var staleCandidates = 0L
+ var fetchSkippedLogResident = 0L
+ var fetchKeyMismatch = 0L
+ var keyLookupRowsDecoded = 0L
+ var metrics = PositionalParquetFetcher.FetchMetrics()
+ val fetcher = new PositionalParquetFetcher(
+ storageConf.newInstance().unwrap(),
+ embeddingCol,
+ vectorSchema,
+ exactQueryVector,
+ ctx.metric,
+ outputSchema,
+ mergedOptions.getOrElse("vector.exact_fetch.offset_index_missing_policy", "fallback").equalsIgnoreCase("fail"),
+ vectorOnlyProjection,
+ VectorIndexOptions.shouldVerifyFetchKeys(mergedOptions.asJava))
+ val keyLocator = new ParquetRecordKeyLocator(storageConf.newInstance().unwrap())
+ while (batches.hasNext) {
+ val (fgId, candidates) = batches.next()
+ batchCount += 1
+ val fileSliceOpt = bFileSliceMap.value.get(fgId)
+ if (fileSliceOpt.isEmpty) {
+ missingSlices += 1
+ throw new HoodieAnalysisException(
+ s"Vector exact search could not resolve file group '$fgId' for ${candidates.size} rerank candidates. " +
+ s"No record-level-index fallback is active in this path.")
+ } else {
+ resolvedSlices += 1
+ val fileSlice = fileSliceOpt.get
+ val logFileCount = fileSlice.getLogFiles.count()
+ val currentBaseInstantTime = fileSlice.getBaseInstantTime
+ val logResident = candidates.filter(candidate =>
+ isLogResidentCandidate(candidate, currentBaseInstantTime, logFileCount > 0))
+ if (logResident.nonEmpty) {
+ fetchSkippedLogResident += logResident.size
+ fetchSkippedLogResidentAcc.add(logResident.size)
+ LOG.warn(
+ s"Vector exact fetch skipped ${logResident.size} log-resident candidate(s) in file group '$fgId'; " +
+ "they remain a result shortfall and the retained candidate pool will refill where candidates remain")
+ }
+ val liveCandidates = candidates.filterNot(logResident.toSet)
+ staleCandidates += liveCandidates.count(candidate =>
+ candidate.getArbiterDecision == VectorIndexArbiter.Decision.STALE)
+ if (liveCandidates.nonEmpty) {
+ val positionalCandidates = liveCandidates.filter { candidate =>
+ val location = candidate.getLocation
+ location != null && location.getPosition >= 0 &&
+ candidate.getArbiterDecision != VectorIndexArbiter.Decision.STALE &&
+ (location.getInstantTime == null || location.getInstantTime == currentBaseInstantTime)
+ }
+ val keyCandidates = liveCandidates.filterNot(positionalCandidates.toSet)
+ val recordKeys = liveCandidates.map(_.getRecordKey).toSet
+ candidateKeys += recordKeys.size
+ if (fileSlice.getBaseFile.isEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact key fallback requires a base Parquet file for file group '$fgId'; " +
+ "the remaining candidates are not classified as log-resident.")
+ }
+ baseFiles += 1
+ logFiles += logFileCount
+ totalSliceBytes += fileSlice.getTotalFileSize
+ val baseFilePath = fileSlice.getBaseFile.get.getPath
+ val locatedByKey = keyLocator.locate(baseFilePath, keyCandidates)
+ keyLookupRowsDecoded += locatedByKey.metrics.rowsDecoded
+ if (locatedByKey.metrics.located != keyCandidates.size) {
+ throw new HoodieAnalysisException(
+ s"Vector exact key fallback located ${locatedByKey.metrics.located} of ${keyCandidates.size} " +
+ s"candidate key(s) in the live base file for file group '$fgId'.")
+ }
+ val firstFetch = fetcher.fetch(baseFilePath, positionalCandidates ++ locatedByKey.candidates)
+ rows ++= firstFetch.rows
+ metrics = metrics.add(firstFetch.metrics)
+ rowsMaterializedAcc.add(firstFetch.metrics.rowsMaterialized)
+ if (firstFetch.retryByKey.nonEmpty) {
+ fetchKeyMismatch += firstFetch.retryByKey.size
+ fetchKeyMismatchAcc.add(firstFetch.retryByKey.size)
+ val retries = keyLocator.locate(baseFilePath, firstFetch.retryByKey)
+ keyLookupRowsDecoded += retries.metrics.rowsDecoded
+ if (retries.metrics.located != firstFetch.retryByKey.size) {
+ throw new HoodieAnalysisException(
+ s"Vector exact fetch could not relocate ${firstFetch.retryByKey.size - retries.metrics.located} " +
+ s"key mismatch(es) in file group '$fgId'.")
+ }
+ val retryFetch = fetcher.fetch(baseFilePath, retries.candidates)
+ if (retryFetch.retryByKey.nonEmpty) {
+ throw new HoodieAnalysisException(
+ s"Vector exact fetch could not self-heal ${retryFetch.retryByKey.size} key mismatch(es) " +
+ s"in file group '$fgId' after key-based relocation.")
+ }
+ rows ++= retryFetch.rows
+ metrics = metrics.add(retryFetch.metrics)
+ rowsMaterializedAcc.add(retryFetch.metrics.rowsMaterialized)
+ }
+ }
+ }
+ }
+ LOG.info(
+ s"[vector_search][stage][exact_read] rerankCandidates=$candidateKeys candidateFileCount=$resolvedSlices " +
+ s"fetchBatches=$batchCount resolvedSlices=$resolvedSlices missingSlices=$missingSlices " +
+ s"rowsDecoded=${metrics.rowsDecoded} rowsMaterialized=${metrics.rowsMaterialized} " +
+ s"keyLookupRowsDecoded=$keyLookupRowsDecoded fetchKeyMismatch=$fetchKeyMismatch " +
+ s"fetchSkippedLogResident=$fetchSkippedLogResident " +
+ s"rowGroupsTotal=${metrics.rowGroupsTotal} rowGroupsSelected=${metrics.rowGroupsSelected} " +
+ s"pagesInSelectedRowGroups=${metrics.pagesInSelectedRowGroups} pagesSelected=${metrics.pagesSelected} " +
+ s"offsetIndexHits=${metrics.offsetIndexHits} offsetIndexMissing=${metrics.offsetIndexMissing} " +
+ s"rangedGets=${metrics.rangedGets} pageBytesFetched=${metrics.pageBytesFetched} " +
+ s"physicalBytesRead=${metrics.physicalBytesRead} footerReads=${metrics.footerReads} footerCacheHits=${metrics.footerCacheHits} " +
+ s"fetchKeyMismatchMetric=${metrics.fetchKeyMismatch} staleCandidates=$staleCandidates " +
+ s"baseFiles=$baseFiles logFiles=$logFiles " +
+ s"sliceBytes=$totalSliceBytes footerMs=${metrics.footerMs} offsetIndexMs=${metrics.offsetIndexMs} " +
+ s"fetchWaitMs=${metrics.fetchWaitMs} decodeMs=${metrics.decodeMs} scoreMs=${metrics.scoreMs} elapsedMs=${elapsedMs(startNs)}")
+ rows.iterator
+ }
+ val logResidentRows: Seq[InternalRow] =
+ if (logResidentCandidates.isEmpty) {
+ Seq.empty
+ } else {
+ val dataSchema = new TableSchemaResolver(ctx.metaClient).getTableSchema(true)
+ val scorer = new VectorExactScorer(vectorSchema, exactQueryVector, ctx.metric)
+ val logFetcher = new LogResidentVectorFetcher(
+ storageConf, ctx.metaClient, dataSchema, ctx.latestInstantTime,
+ embeddingCol, outputSchema, scorer)
+ val fetched = logResidentCandidates
+ .groupBy(candidate =>
+ (Option(candidate.getLocation.getPartitionPath).getOrElse(""), candidate.getLocation.getFileId))
+ .flatMap { case ((partitionPath, fgId), candidates) =>
+ logFetcher.fetchSlice(partitionPath, fileSliceMap(fgId), candidates.map(_.getRecordKey).toSet)
+ }.toSeq
+ LOG.info(
+ s"[vector_search][exact][log_resident] candidates=${logResidentCandidates.size} " +
+ s"materialized=${fetched.size} " +
+ s"fileGroups=${logResidentCandidates.map(_.getLocation.getFileId).distinct.size}")
+ fetched
+ }
+ val combinedRows: RDD[InternalRow] =
+ if (logResidentRows.isEmpty) exactRows
+ else exactRows.union(spark.sparkContext.parallelize(logResidentRows))
+ LOG.info(s"[vector_search][plan] exact candidate DF planned in ${elapsedMs(planStartNs)} ms")
+ sparkAdapter.getUnsafeUtils.createDataFrameFromRDD(spark, combinedRows, outputSchema)
+ }
+ }
+ }
+
+ private def preResolveFileSlices(
+ metadataTable: HoodieTableMetadata,
+ metaClient: HoodieTableMetaClient,
+ timeline: HoodieTimeline,
+ latestInstantTime: String,
+ candidatesByPartition: Map[String, Set[String]]): Map[String, FileSlice] = {
+ val startNs = System.nanoTime()
+ val fsView = new HoodieTableFileSystemView(metadataTable, metaClient, timeline)
+ try {
+ val result = scala.collection.mutable.Map.empty[String, FileSlice]
+ var baseFiles = 0
+ var logFiles = 0L
+ var totalSliceBytes = 0L
+ candidatesByPartition.foreach { case (partitionPath, candidateFgIds) =>
+ val sliceIter = fsView.getLatestFileSlicesBeforeOrOn(partitionPath, latestInstantTime, true).iterator()
+ while (sliceIter.hasNext) {
+ val slice = sliceIter.next()
+ val fgId = slice.getFileGroupId.getFileId
+ if (candidateFgIds.contains(fgId)) {
+ result(fgId) = slice
Review Comment:
**Fixed**. Slices are keyed by `HoodieFileGroupId(partitionPath, fileId)` (`fileGroupIdOf`), so reused bucket file ids in different partitions don't collide.
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/PositionalParquetFetcher.scala:
##########
@@ -0,0 +1,587 @@
+/*
+ * 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.spark.sql.hudi.analysis
+
+import org.apache.hudi.common.index.vector.{ScoredVectorPostingMatch, VectorIndexMdtSearchUtils, VectorPostingMatch}
+import org.apache.hudi.common.model.HoodieRecord
+import org.apache.hudi.common.schema.HoodieSchema
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileSystem, Path}
+import org.apache.parquet.column.page.PageReadStore
+import org.apache.parquet.hadoop.ParquetFileReader
+import org.apache.parquet.hadoop.metadata.{BlockMetaData, ColumnChunkMetaData, ParquetMetadata}
+import org.apache.parquet.internal.column.columnindex.OffsetIndex
+import org.apache.parquet.internal.filter2.columnindex.RowRanges
+import org.apache.parquet.io.{ColumnIOFactory, RecordReader}
+import org.apache.parquet.io.api.{Binary, Converter, GroupConverter, PrimitiveConverter, RecordMaterializer}
+import org.apache.parquet.schema.MessageType
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction.DistanceMetric
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+import org.slf4j.LoggerFactory
+
+import java.util.PrimitiveIterator
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+/**
+ * Fetches exact vector-search rerank candidates from Parquet by physical row position.
+ *
+ * The fetcher uses Parquet page indexes to turn candidate row positions into row ranges, so
+ * the reader can skip non-candidate pages instead of scanning whole Hudi file slices.
+ */
+private[analysis] final class PositionalParquetFetcher(
+ conf: Configuration,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Double],
+ metric: DistanceMetric.Value,
+ outputSchema: StructType,
+ failOnMissingOffsetIndex: Boolean,
+ vectorOnlyProjection: Boolean,
+ verifyKeys: Boolean) extends Serializable {
+
+ import PositionalParquetFetcher._
+
+ private val distanceOrdinal = outputSchema.fieldIndex(HoodieVectorSearchPlanBuilder.DISTANCE_COL)
+ private val recordKeyOrdinal = fieldOrdinal(HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ private val partitionPathOrdinal = fieldOrdinal(HoodieRecord.PARTITION_PATH_METADATA_FIELD)
+ private val fileGroupIdOrdinal = fieldOrdinal(HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL)
+ private val scorer = new VectorExactScorer(vectorSchema, queryVector, metric)
+ private var scoreMs: Long = 0L
+
+ def fetch(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch]): FetchResult = {
+ val startNs = System.nanoTime()
+ if (candidates.isEmpty) {
+ FetchResult(Iterator.empty, Seq.empty, FetchMetrics(elapsedMs = elapsedMs(startNs)))
+ } else {
+ fetchNonEmpty(filePath, candidates, startNs)
+ }
+ }
+
+ private def fetchNonEmpty(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch],
+ startNs: Long): FetchResult = {
+ val path = new Path(filePath)
+ configureRandomRead(conf)
+ val fs = path.getFileSystem(conf)
+ val fileLen = fs.getFileStatus(path).getLen
+ val fileCacheKey = s"$filePath#$fileLen"
+ val stats = Option(FileSystem.getStatistics(path.toUri.getScheme, fs.getClass))
+ val physicalBefore = stats.map(_.getBytesRead).getOrElse(0L)
+
+ val footerStartNs = System.nanoTime()
+ val footerResult = footer(conf, path, fileCacheKey)
+ val footerMs = elapsedMs(footerStartNs)
+ val metadata = footerResult.metadata
+ val fileSchema = metadata.getFileMetaData.getSchema
+ val requestedParquetSchema = requestedSchema(fileSchema)
+
+ val rowGroups = metadata.getBlocks.asScala.toIndexedSeq
+ val rowGroupStarts = cumulativeRowGroupStarts(rowGroups)
+ val byRowGroup: Map[Int, Seq[ScoredVectorPostingMatch]] =
+ candidates.groupBy(candidate => rowGroupOrdinal(rowGroupStarts, rowGroups, rowPosition(candidate)))
+
+ val rows = mutable.ArrayBuffer.empty[InternalRow]
+ val keyMismatches = mutable.ArrayBuffer.empty[ScoredVectorPostingMatch]
+ var rowGroupsSelected = 0
+ var pagesInSelectedRowGroups = 0L
+ var pagesSelected = 0L
+ var offsetIndexHits = 0L
+ var offsetIndexMissing = 0L
+ var pageBytesFetched = 0L
+ var rangedGets = 0L
+ var rowsDecoded = 0L
+ var rowsMaterialized = 0L
+ var offsetIndexMs = 0L
+ var decodeMs = 0L
+ scoreMs = 0L
+
+ val reader = new ParquetFileReader(conf, path, metadata)
+ try {
+ reader.setRequestedSchema(requestedParquetSchema)
+ byRowGroup.toSeq.sortBy(_._1).foreach { case (rgOrdinal, rgCandidates) =>
+ val block = rowGroups(rgOrdinal)
+ val rowGroupStart = rowGroupStarts(rgOrdinal)
+ val embeddingChunk = findColumn(block, embeddingCol)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet file '$filePath' is missing vector column '$embeddingCol'."))
+
+ val offsetStartNs = System.nanoTime()
+ val embeddingOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, embeddingChunk)
+ offsetIndexMs += elapsedMs(offsetStartNs)
+ val relativeRows = rgCandidates.map(candidate => rowPosition(candidate) - rowGroupStart).sorted
+ val candidateRows = relativeRows.toSet
+ val candidateByRelativeRow = rgCandidates
+ .map(candidate => rowPosition(candidate) - rowGroupStart -> candidate)
+ .toMap
+
+ if (embeddingOffsetIndex == null && failOnMissingOffsetIndex) {
+ offsetIndexMissing += 1
+ throw new IllegalStateException(
+ s"Parquet file '$filePath' has no OffsetIndex for vector column '$embeddingCol'; positional exact fetch cannot honor pages.")
+ }
+
+ rowGroupsSelected += 1
+ val (pageStore, rowRanges) = if (embeddingOffsetIndex == null) {
+ offsetIndexMissing += 1
+ LOG.warn(
+ "Parquet file {} rowGroup={} has no OffsetIndex for vector column {}; falling back to row-group fetch.",
+ filePath,
+ Int.box(rgOrdinal),
+ embeddingCol)
+ pagesInSelectedRowGroups += 0L
+ pagesSelected += 0L
+ rangedGets += 1L
+ pageBytesFetched += requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .map(_.getTotalSize)
+ .sum
+ (reader.readRowGroup(rgOrdinal), RowRanges.createSingle(block.getRowCount))
+ } else {
+ offsetIndexHits += 1
+ pagesInSelectedRowGroups += embeddingOffsetIndex.getPageCount
+ val selectedVectorPages = selectedPages(embeddingOffsetIndex, block.getRowCount, relativeRows)
+ pagesSelected += selectedVectorPages.size
+ val selectedRanges = requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .flatMap { chunk =>
+ val oiStartNs = System.nanoTime()
+ val chunkOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, chunk)
+ offsetIndexMs += elapsedMs(oiStartNs)
+ if (chunkOffsetIndex == null) {
+ offsetIndexMissing += 1
+ Seq.empty
+ } else {
+ offsetIndexHits += 1
+ selectedPages(chunkOffsetIndex, block.getRowCount, relativeRows).map { pageOrdinal =>
+ val start = chunkOffsetIndex.getOffset(pageOrdinal)
+ PageRange(start, start + chunkOffsetIndex.getCompressedPageSize(pageOrdinal))
+ }
+ }
+ }
+ val mergedRanges = mergeRanges(selectedRanges.sortBy(_.start).toSeq)
+ rangedGets += mergedRanges.size
+ pageBytesFetched += mergedRanges.iterator.map(range => range.end - range.start).sum
+ val ranges = RowRanges.create(block.getRowCount, primitiveIntIterator(selectedVectorPages), embeddingOffsetIndex)
+ (reader.readFilteredRowGroup(rgOrdinal, ranges), ranges)
+ }
+
+ val decodeStartNs = System.nanoTime()
+ val recordReader = parquetRecordReader(fileSchema, requestedParquetSchema, pageStore)
+ val rowIterator = rowRanges.iterator()
+ while (rowIterator.hasNext) {
+ val relativeRow = rowIterator.nextLong()
+ val row = recordReader.read()
+ rowsDecoded += 1
+ candidateByRelativeRow.get(relativeRow).foreach { candidate =>
+ if (verifyKeys && decodedRecordKey(row) != candidate.getRecordKey) {
+ keyMismatches += candidate
+ } else {
+ fillSyntheticFields(row, candidate)
+ rows += row
+ rowsMaterialized += 1
+ }
+ }
+ }
+ decodeMs += elapsedMs(decodeStartNs)
+ }
+ } finally {
+ reader.close()
+ }
+
+ val physicalAfter = stats.map(_.getBytesRead).getOrElse(physicalBefore)
+ val metrics = FetchMetrics(
+ rerankCandidates = candidates.size,
+ rowGroupsTotal = rowGroups.size,
+ rowGroupsSelected = rowGroupsSelected,
+ pagesInSelectedRowGroups = pagesInSelectedRowGroups,
+ pagesSelected = pagesSelected,
+ offsetIndexHits = offsetIndexHits,
+ offsetIndexMissing = offsetIndexMissing,
+ rangedGets = rangedGets,
+ pageBytesFetched = pageBytesFetched,
+ rowsDecoded = rowsDecoded,
+ rowsMaterialized = rowsMaterialized,
+ fetchKeyMismatch = keyMismatches.size,
+ physicalBytesRead = math.max(0L, physicalAfter - physicalBefore),
+ footerReads = if (footerResult.cacheHit) 0 else 1,
+ footerCacheHits = if (footerResult.cacheHit) 1 else 0,
+ footerMs = footerMs,
+ offsetIndexMs = offsetIndexMs,
+ fetchWaitMs = 0L,
+ decodeMs = decodeMs,
+ scoreMs = scoreMs,
+ elapsedMs = elapsedMs(startNs))
+ FetchResult(rows.iterator, keyMismatches.toSeq, metrics)
+ }
+
+ private def parquetRecordReader(
+ fileSchema: MessageType,
+ requestedParquetSchema: MessageType,
+ pageStore: PageReadStore): RecordReader[InternalRow] = {
+ val materializer = new InternalRowMaterializer(requestedParquetSchema, outputSchema, embeddingCol, distanceOrdinal)
+ new ColumnIOFactory().getColumnIO(requestedParquetSchema, fileSchema)
+ .getRecordReader(pageStore, materializer)
+ }
+
+ private def requestedSchema(fileSchema: MessageType): MessageType = {
+ val parquetFields =
+ if (vectorOnlyProjection) {
+ mutable.ArrayBuffer.empty[org.apache.parquet.schema.Type]
+ } else {
+ outputSchema.fields
+ .filterNot(field => isSyntheticOutputField(field.name))
+ .map(field => fieldByName(fileSchema, field.name))
+ .toBuffer[org.apache.parquet.schema.Type]
+ }
+ if ((verifyKeys || vectorOnlyProjection) && !parquetFields.exists(_.getName == HoodieRecord.RECORD_KEY_METADATA_FIELD)) {
+ parquetFields += fieldByName(fileSchema, HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ }
+ if (!parquetFields.exists(_.getName == embeddingCol)) {
+ parquetFields += fieldByName(fileSchema, embeddingCol)
+ }
+ new MessageType(fileSchema.getName, parquetFields.asJava)
+ }
+
+ private def isSyntheticOutputField(name: String): Boolean =
+ name == HoodieVectorSearchPlanBuilder.DISTANCE_COL ||
+ name == HoodieRecord.RECORD_KEY_METADATA_FIELD ||
+ name == HoodieRecord.PARTITION_PATH_METADATA_FIELD ||
+ name == HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL
+
+ private def fieldByName(schema: MessageType, name: String): org.apache.parquet.schema.Type =
+ schema.getFields.asScala.find(_.getName == name)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet schema is missing requested field '$name'."))
+
+ private def rowPosition(candidate: ScoredVectorPostingMatch): Long = {
+ val location = candidate.getLocation
+ if (location == null) {
+ throw new IllegalArgumentException(
+ s"Vector exact positional fetch requires a row-position locator for record ${candidate.getRecordKey}.")
+ }
+ location.getPosition
+ }
+
+ private def decodedRecordKey(row: InternalRow): String =
+ recordKeyOrdinal.flatMap(ordinal => Option(row.getUTF8String(ordinal))).map(_.toString).orNull
+
+ private def fillSyntheticFields(row: InternalRow, candidate: ScoredVectorPostingMatch): Unit = {
+ val liveLocation = Option(candidate.getLocation)
+ recordKeyOrdinal.foreach(row.update(_, UTF8String.fromString(candidate.getRecordKey)))
+ partitionPathOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getPartitionPath).orElse(Option(candidate.getPartitionPath)).getOrElse(""))))
+ fileGroupIdOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getFileId).orElse(Option(candidate.getFileGroupId)).getOrElse(""))))
+ }
+
+ private def fieldOrdinal(name: String): Option[Int] =
+ outputSchema.fields.indexWhere(_.name == name) match {
+ case -1 => None
+ case ordinal => Some(ordinal)
+ }
+
+ private def scoreVector(bytes: Array[Byte]): Double = scorer.scoreBytes(bytes)
+
+ private final class InternalRowMaterializer(
+ requestedParquetSchema: MessageType,
+ outputSchema: StructType,
+ embeddingCol: String,
+ distanceOrdinal: Int) extends RecordMaterializer[InternalRow] {
+
+ private val values = new Array[Any](outputSchema.length)
+ private var current: InternalRow = _
+ private var vectorBytes: Array[Byte] = _
+
+ private val root = new GroupConverter {
+ private val converters = requestedParquetSchema.getFields.asScala.map(converterFor).toIndexedSeq
+
+ override def getConverter(fieldIndex: Int): Converter =
+ converters(fieldIndex)
+
+ override def start(): Unit = {
+ var idx = 0
+ while (idx < values.length) {
+ values(idx) = null
+ idx += 1
+ }
+ vectorBytes = null
+ }
+
+ override def end(): Unit = {
+ if (vectorBytes == null) {
+ throw new IllegalStateException(s"Vector column '$embeddingCol' was not decoded for positional exact fetch.")
+ }
+ val scoreStartNs = System.nanoTime()
+ values(distanceOrdinal) = scoreVector(vectorBytes)
+ scoreMs += elapsedMs(scoreStartNs)
+ current = new GenericInternalRow(values.clone())
+ }
+ }
+
+ override def getCurrentRecord: InternalRow = current
+
+ override def getRootConverter: GroupConverter = root
+
+ private def converterFor(field: org.apache.parquet.schema.Type): PrimitiveConverter = {
+ val name = field.getName
+ new PrimitiveConverter {
+ override def addBinary(value: Binary): Unit = {
+ if (name == embeddingCol) {
+ vectorBytes = value.getBytes
+ } else {
+ setOutputValue(name, value.getBytes)
+ }
+ }
+
+ override def addInt(value: Int): Unit = setOutputValue(name, value)
+
+ override def addLong(value: Long): Unit = setOutputValue(name, value)
+
+ override def addFloat(value: Float): Unit = setOutputValue(name, value)
+
+ override def addDouble(value: Double): Unit = setOutputValue(name, value)
+
+ override def addBoolean(value: Boolean): Unit = setOutputValue(name, value)
+ }
+ }
+
+ private def setOutputValue(name: String, raw: Any): Unit = {
+ if (name != embeddingCol && name != HoodieVectorSearchPlanBuilder.DISTANCE_COL) {
+ val ordinal = outputSchema.fieldIndex(name)
+ values(ordinal) = outputSchema.fields(ordinal).dataType match {
+ case StringType => raw match {
+ case bytes: Array[Byte] => UTF8String.fromBytes(bytes)
+ case other => UTF8String.fromString(String.valueOf(other))
+ }
+ case BinaryType => raw.asInstanceOf[Array[Byte]]
+ case IntegerType => raw.asInstanceOf[Int]
+ case LongType => raw.asInstanceOf[Long]
+ case FloatType => raw.asInstanceOf[Float]
+ case DoubleType => raw.asInstanceOf[Double]
+ case BooleanType => raw.asInstanceOf[Boolean]
+ case ByteType => raw.asInstanceOf[Int].toByte
+ case ShortType => raw.asInstanceOf[Int].toShort
+ case other => throw new UnsupportedOperationException(
+ s"Unsupported exact positional fetch output type for field '$name': $other")
+ }
+ }
+ }
+ }
+}
+
+private[analysis] object PositionalParquetFetcher {
+ private val LOG = LoggerFactory.getLogger(getClass)
+ private val FOOTER_CACHE = new ConcurrentHashMap[String, ParquetMetadata]()
Review Comment:
**Fixed** `3cde5283ea`. Footer/offset-index caches are bounded Caffeine caches.
##########
hudi-spark-datasource/hudi-spark-common/src/main/scala/org/apache/spark/sql/hudi/analysis/PositionalParquetFetcher.scala:
##########
@@ -0,0 +1,587 @@
+/*
+ * 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.spark.sql.hudi.analysis
+
+import org.apache.hudi.common.index.vector.{ScoredVectorPostingMatch, VectorIndexMdtSearchUtils, VectorPostingMatch}
+import org.apache.hudi.common.model.HoodieRecord
+import org.apache.hudi.common.schema.HoodieSchema
+
+import org.apache.hadoop.conf.Configuration
+import org.apache.hadoop.fs.{FileSystem, Path}
+import org.apache.parquet.column.page.PageReadStore
+import org.apache.parquet.hadoop.ParquetFileReader
+import org.apache.parquet.hadoop.metadata.{BlockMetaData, ColumnChunkMetaData, ParquetMetadata}
+import org.apache.parquet.internal.column.columnindex.OffsetIndex
+import org.apache.parquet.internal.filter2.columnindex.RowRanges
+import org.apache.parquet.io.{ColumnIOFactory, RecordReader}
+import org.apache.parquet.io.api.{Binary, Converter, GroupConverter, PrimitiveConverter, RecordMaterializer}
+import org.apache.parquet.schema.MessageType
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.GenericInternalRow
+import org.apache.spark.sql.catalyst.plans.logical.HoodieVectorSearchTableValuedFunction.DistanceMetric
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
+import org.slf4j.LoggerFactory
+
+import java.util.PrimitiveIterator
+import java.util.concurrent.ConcurrentHashMap
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+
+/**
+ * Fetches exact vector-search rerank candidates from Parquet by physical row position.
+ *
+ * The fetcher uses Parquet page indexes to turn candidate row positions into row ranges, so
+ * the reader can skip non-candidate pages instead of scanning whole Hudi file slices.
+ */
+private[analysis] final class PositionalParquetFetcher(
+ conf: Configuration,
+ embeddingCol: String,
+ vectorSchema: HoodieSchema.Vector,
+ queryVector: Array[Double],
+ metric: DistanceMetric.Value,
+ outputSchema: StructType,
+ failOnMissingOffsetIndex: Boolean,
+ vectorOnlyProjection: Boolean,
+ verifyKeys: Boolean) extends Serializable {
+
+ import PositionalParquetFetcher._
+
+ private val distanceOrdinal = outputSchema.fieldIndex(HoodieVectorSearchPlanBuilder.DISTANCE_COL)
+ private val recordKeyOrdinal = fieldOrdinal(HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ private val partitionPathOrdinal = fieldOrdinal(HoodieRecord.PARTITION_PATH_METADATA_FIELD)
+ private val fileGroupIdOrdinal = fieldOrdinal(HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL)
+ private val scorer = new VectorExactScorer(vectorSchema, queryVector, metric)
+ private var scoreMs: Long = 0L
+
+ def fetch(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch]): FetchResult = {
+ val startNs = System.nanoTime()
+ if (candidates.isEmpty) {
+ FetchResult(Iterator.empty, Seq.empty, FetchMetrics(elapsedMs = elapsedMs(startNs)))
+ } else {
+ fetchNonEmpty(filePath, candidates, startNs)
+ }
+ }
+
+ private def fetchNonEmpty(
+ filePath: String,
+ candidates: Seq[ScoredVectorPostingMatch],
+ startNs: Long): FetchResult = {
+ val path = new Path(filePath)
+ configureRandomRead(conf)
+ val fs = path.getFileSystem(conf)
+ val fileLen = fs.getFileStatus(path).getLen
+ val fileCacheKey = s"$filePath#$fileLen"
+ val stats = Option(FileSystem.getStatistics(path.toUri.getScheme, fs.getClass))
+ val physicalBefore = stats.map(_.getBytesRead).getOrElse(0L)
+
+ val footerStartNs = System.nanoTime()
+ val footerResult = footer(conf, path, fileCacheKey)
+ val footerMs = elapsedMs(footerStartNs)
+ val metadata = footerResult.metadata
+ val fileSchema = metadata.getFileMetaData.getSchema
+ val requestedParquetSchema = requestedSchema(fileSchema)
+
+ val rowGroups = metadata.getBlocks.asScala.toIndexedSeq
+ val rowGroupStarts = cumulativeRowGroupStarts(rowGroups)
+ val byRowGroup: Map[Int, Seq[ScoredVectorPostingMatch]] =
+ candidates.groupBy(candidate => rowGroupOrdinal(rowGroupStarts, rowGroups, rowPosition(candidate)))
+
+ val rows = mutable.ArrayBuffer.empty[InternalRow]
+ val keyMismatches = mutable.ArrayBuffer.empty[ScoredVectorPostingMatch]
+ var rowGroupsSelected = 0
+ var pagesInSelectedRowGroups = 0L
+ var pagesSelected = 0L
+ var offsetIndexHits = 0L
+ var offsetIndexMissing = 0L
+ var pageBytesFetched = 0L
+ var rangedGets = 0L
+ var rowsDecoded = 0L
+ var rowsMaterialized = 0L
+ var offsetIndexMs = 0L
+ var decodeMs = 0L
+ scoreMs = 0L
+
+ val reader = new ParquetFileReader(conf, path, metadata)
+ try {
+ reader.setRequestedSchema(requestedParquetSchema)
+ byRowGroup.toSeq.sortBy(_._1).foreach { case (rgOrdinal, rgCandidates) =>
+ val block = rowGroups(rgOrdinal)
+ val rowGroupStart = rowGroupStarts(rgOrdinal)
+ val embeddingChunk = findColumn(block, embeddingCol)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet file '$filePath' is missing vector column '$embeddingCol'."))
+
+ val offsetStartNs = System.nanoTime()
+ val embeddingOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, embeddingChunk)
+ offsetIndexMs += elapsedMs(offsetStartNs)
+ val relativeRows = rgCandidates.map(candidate => rowPosition(candidate) - rowGroupStart).sorted
+ val candidateRows = relativeRows.toSet
+ val candidateByRelativeRow = rgCandidates
+ .map(candidate => rowPosition(candidate) - rowGroupStart -> candidate)
+ .toMap
+
+ if (embeddingOffsetIndex == null && failOnMissingOffsetIndex) {
+ offsetIndexMissing += 1
+ throw new IllegalStateException(
+ s"Parquet file '$filePath' has no OffsetIndex for vector column '$embeddingCol'; positional exact fetch cannot honor pages.")
+ }
+
+ rowGroupsSelected += 1
+ val (pageStore, rowRanges) = if (embeddingOffsetIndex == null) {
+ offsetIndexMissing += 1
+ LOG.warn(
+ "Parquet file {} rowGroup={} has no OffsetIndex for vector column {}; falling back to row-group fetch.",
+ filePath,
+ Int.box(rgOrdinal),
+ embeddingCol)
+ pagesInSelectedRowGroups += 0L
+ pagesSelected += 0L
+ rangedGets += 1L
+ pageBytesFetched += requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .map(_.getTotalSize)
+ .sum
+ (reader.readRowGroup(rgOrdinal), RowRanges.createSingle(block.getRowCount))
+ } else {
+ offsetIndexHits += 1
+ pagesInSelectedRowGroups += embeddingOffsetIndex.getPageCount
+ val selectedVectorPages = selectedPages(embeddingOffsetIndex, block.getRowCount, relativeRows)
+ pagesSelected += selectedVectorPages.size
+ val selectedRanges = requestedParquetSchema.getFields.asScala
+ .flatMap(field => findColumn(block, field.getName))
+ .flatMap { chunk =>
+ val oiStartNs = System.nanoTime()
+ val chunkOffsetIndex = offsetIndex(reader, fileCacheKey, rgOrdinal, chunk)
+ offsetIndexMs += elapsedMs(oiStartNs)
+ if (chunkOffsetIndex == null) {
+ offsetIndexMissing += 1
+ Seq.empty
+ } else {
+ offsetIndexHits += 1
+ selectedPages(chunkOffsetIndex, block.getRowCount, relativeRows).map { pageOrdinal =>
+ val start = chunkOffsetIndex.getOffset(pageOrdinal)
+ PageRange(start, start + chunkOffsetIndex.getCompressedPageSize(pageOrdinal))
+ }
+ }
+ }
+ val mergedRanges = mergeRanges(selectedRanges.sortBy(_.start).toSeq)
+ rangedGets += mergedRanges.size
+ pageBytesFetched += mergedRanges.iterator.map(range => range.end - range.start).sum
+ val ranges = RowRanges.create(block.getRowCount, primitiveIntIterator(selectedVectorPages), embeddingOffsetIndex)
+ (reader.readFilteredRowGroup(rgOrdinal, ranges), ranges)
+ }
+
+ val decodeStartNs = System.nanoTime()
+ val recordReader = parquetRecordReader(fileSchema, requestedParquetSchema, pageStore)
+ val rowIterator = rowRanges.iterator()
+ while (rowIterator.hasNext) {
+ val relativeRow = rowIterator.nextLong()
+ val row = recordReader.read()
+ rowsDecoded += 1
+ candidateByRelativeRow.get(relativeRow).foreach { candidate =>
+ if (verifyKeys && decodedRecordKey(row) != candidate.getRecordKey) {
+ keyMismatches += candidate
+ } else {
+ fillSyntheticFields(row, candidate)
+ rows += row
+ rowsMaterialized += 1
+ }
+ }
+ }
+ decodeMs += elapsedMs(decodeStartNs)
+ }
+ } finally {
+ reader.close()
+ }
+
+ val physicalAfter = stats.map(_.getBytesRead).getOrElse(physicalBefore)
+ val metrics = FetchMetrics(
+ rerankCandidates = candidates.size,
+ rowGroupsTotal = rowGroups.size,
+ rowGroupsSelected = rowGroupsSelected,
+ pagesInSelectedRowGroups = pagesInSelectedRowGroups,
+ pagesSelected = pagesSelected,
+ offsetIndexHits = offsetIndexHits,
+ offsetIndexMissing = offsetIndexMissing,
+ rangedGets = rangedGets,
+ pageBytesFetched = pageBytesFetched,
+ rowsDecoded = rowsDecoded,
+ rowsMaterialized = rowsMaterialized,
+ fetchKeyMismatch = keyMismatches.size,
+ physicalBytesRead = math.max(0L, physicalAfter - physicalBefore),
+ footerReads = if (footerResult.cacheHit) 0 else 1,
+ footerCacheHits = if (footerResult.cacheHit) 1 else 0,
+ footerMs = footerMs,
+ offsetIndexMs = offsetIndexMs,
+ fetchWaitMs = 0L,
+ decodeMs = decodeMs,
+ scoreMs = scoreMs,
+ elapsedMs = elapsedMs(startNs))
+ FetchResult(rows.iterator, keyMismatches.toSeq, metrics)
+ }
+
+ private def parquetRecordReader(
+ fileSchema: MessageType,
+ requestedParquetSchema: MessageType,
+ pageStore: PageReadStore): RecordReader[InternalRow] = {
+ val materializer = new InternalRowMaterializer(requestedParquetSchema, outputSchema, embeddingCol, distanceOrdinal)
+ new ColumnIOFactory().getColumnIO(requestedParquetSchema, fileSchema)
+ .getRecordReader(pageStore, materializer)
+ }
+
+ private def requestedSchema(fileSchema: MessageType): MessageType = {
+ val parquetFields =
+ if (vectorOnlyProjection) {
+ mutable.ArrayBuffer.empty[org.apache.parquet.schema.Type]
+ } else {
+ outputSchema.fields
+ .filterNot(field => isSyntheticOutputField(field.name))
+ .map(field => fieldByName(fileSchema, field.name))
+ .toBuffer[org.apache.parquet.schema.Type]
+ }
+ if ((verifyKeys || vectorOnlyProjection) && !parquetFields.exists(_.getName == HoodieRecord.RECORD_KEY_METADATA_FIELD)) {
+ parquetFields += fieldByName(fileSchema, HoodieRecord.RECORD_KEY_METADATA_FIELD)
+ }
+ if (!parquetFields.exists(_.getName == embeddingCol)) {
+ parquetFields += fieldByName(fileSchema, embeddingCol)
+ }
+ new MessageType(fileSchema.getName, parquetFields.asJava)
+ }
+
+ private def isSyntheticOutputField(name: String): Boolean =
+ name == HoodieVectorSearchPlanBuilder.DISTANCE_COL ||
+ name == HoodieRecord.RECORD_KEY_METADATA_FIELD ||
+ name == HoodieRecord.PARTITION_PATH_METADATA_FIELD ||
+ name == HoodieVectorSearchPlanBuilder.FILE_GROUP_ID_COL
+
+ private def fieldByName(schema: MessageType, name: String): org.apache.parquet.schema.Type =
+ schema.getFields.asScala.find(_.getName == name)
+ .getOrElse(throw new IllegalArgumentException(s"Parquet schema is missing requested field '$name'."))
+
+ private def rowPosition(candidate: ScoredVectorPostingMatch): Long = {
+ val location = candidate.getLocation
+ if (location == null) {
+ throw new IllegalArgumentException(
+ s"Vector exact positional fetch requires a row-position locator for record ${candidate.getRecordKey}.")
+ }
+ location.getPosition
+ }
+
+ private def decodedRecordKey(row: InternalRow): String =
+ recordKeyOrdinal.flatMap(ordinal => Option(row.getUTF8String(ordinal))).map(_.toString).orNull
+
+ private def fillSyntheticFields(row: InternalRow, candidate: ScoredVectorPostingMatch): Unit = {
+ val liveLocation = Option(candidate.getLocation)
+ recordKeyOrdinal.foreach(row.update(_, UTF8String.fromString(candidate.getRecordKey)))
+ partitionPathOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getPartitionPath).orElse(Option(candidate.getPartitionPath)).getOrElse(""))))
+ fileGroupIdOrdinal.foreach(row.update(_, UTF8String.fromString(
+ liveLocation.map(_.getFileId).orElse(Option(candidate.getFileGroupId)).getOrElse(""))))
+ }
+
+ private def fieldOrdinal(name: String): Option[Int] =
+ outputSchema.fields.indexWhere(_.name == name) match {
+ case -1 => None
+ case ordinal => Some(ordinal)
+ }
+
+ private def scoreVector(bytes: Array[Byte]): Double = scorer.scoreBytes(bytes)
+
+ private final class InternalRowMaterializer(
+ requestedParquetSchema: MessageType,
+ outputSchema: StructType,
+ embeddingCol: String,
+ distanceOrdinal: Int) extends RecordMaterializer[InternalRow] {
+
+ private val values = new Array[Any](outputSchema.length)
+ private var current: InternalRow = _
+ private var vectorBytes: Array[Byte] = _
+
+ private val root = new GroupConverter {
+ private val converters = requestedParquetSchema.getFields.asScala.map(converterFor).toIndexedSeq
+
+ override def getConverter(fieldIndex: Int): Converter =
+ converters(fieldIndex)
+
+ override def start(): Unit = {
+ var idx = 0
+ while (idx < values.length) {
+ values(idx) = null
+ idx += 1
+ }
+ vectorBytes = null
+ }
+
+ override def end(): Unit = {
+ if (vectorBytes == null) {
+ throw new IllegalStateException(s"Vector column '$embeddingCol' was not decoded for positional exact fetch.")
+ }
+ val scoreStartNs = System.nanoTime()
+ values(distanceOrdinal) = scoreVector(vectorBytes)
+ scoreMs += elapsedMs(scoreStartNs)
+ current = new GenericInternalRow(values.clone())
+ }
+ }
+
+ override def getCurrentRecord: InternalRow = current
+
+ override def getRootConverter: GroupConverter = root
+
+ private def converterFor(field: org.apache.parquet.schema.Type): PrimitiveConverter = {
+ val name = field.getName
+ new PrimitiveConverter {
+ override def addBinary(value: Binary): Unit = {
+ if (name == embeddingCol) {
+ vectorBytes = value.getBytes
+ } else {
+ setOutputValue(name, value.getBytes)
+ }
+ }
+
+ override def addInt(value: Int): Unit = setOutputValue(name, value)
+
+ override def addLong(value: Long): Unit = setOutputValue(name, value)
+
+ override def addFloat(value: Float): Unit = setOutputValue(name, value)
+
+ override def addDouble(value: Double): Unit = setOutputValue(name, value)
+
+ override def addBoolean(value: Boolean): Unit = setOutputValue(name, value)
+ }
+ }
+
+ private def setOutputValue(name: String, raw: Any): Unit = {
+ if (name != embeddingCol && name != HoodieVectorSearchPlanBuilder.DISTANCE_COL) {
+ val ordinal = outputSchema.fieldIndex(name)
+ values(ordinal) = outputSchema.fields(ordinal).dataType match {
+ case StringType => raw match {
+ case bytes: Array[Byte] => UTF8String.fromBytes(bytes)
+ case other => UTF8String.fromString(String.valueOf(other))
+ }
+ case BinaryType => raw.asInstanceOf[Array[Byte]]
+ case IntegerType => raw.asInstanceOf[Int]
+ case LongType => raw.asInstanceOf[Long]
+ case FloatType => raw.asInstanceOf[Float]
+ case DoubleType => raw.asInstanceOf[Double]
+ case BooleanType => raw.asInstanceOf[Boolean]
+ case ByteType => raw.asInstanceOf[Int].toByte
+ case ShortType => raw.asInstanceOf[Int].toShort
+ case other => throw new UnsupportedOperationException(
Review Comment:
**Fixed** `d685311d5f`. The positional fetcher now decodes through Spark's Parquet read support (`sparkAdapter.createParquetReadSupport`) with page-level row-range selection; the log path uses `sparkAdapter.createAvroDeserializer`. Timestamp, date, decimal, struct, array and map are covered for COW and MOR by `testExactSearchReturnsNonPrimitiveColumns`.
--
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: commits-unsubscribe@hudi.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org