feat(index): add RaBitQ encoder and multibit-aware scorer contract - #19318
feat(index): add RaBitQ encoder and multibit-aware scorer contract#19318chrevanthreddy wants to merge 6 commits into
Conversation
c4189f9 to
0682b18
Compare
hudi-agent
left a comment
There was a problem hiding this comment.
Thanks for the contribution! This PR adds an engine-neutral RaBitQ encoder and quantized-vector representation, metric-aware multibit scorer contracts (MetricQueryState, RaBitQByteLutScorer, RaBitQPlaneKernel, VectorQueryPlanes) with posting-block build/view primitives, and the corresponding vector-index metadata-table key/payload/partition plumbing, plus the accompanying test suite. No correctness issues found. A few style/readability suggestions in the inline comments. Please take a look, and this should be ready for a Hudi committer or PMC member to take it from here. One naming inconsistency worth fixing in the new test files. a few naming and readability suggestions below, most notably the Object-typed vector index field that breaks the specific-typed pattern of every other payload field, and a factory overload that silently discards parameters.
cc @yihua
|
|
||
| private static final int DIM = 96; | ||
| private static final int TRIALS = 200; | ||
| private static final int BITS = 4; |
There was a problem hiding this comment.
🤖 nit: could you rename this to TestMetricEstimatorIdentity? Every other test class in this PR and the surrounding package uses the Test* prefix convention (TestRaBitQEncoder, TestMetricQueryStateRotation, etc.) — the *Test suffix stands out as an outlier.
| protected HoodieMetadataColumnStats columnStatMetadata = null; | ||
| protected HoodieRecordIndexInfo recordIndexMetadata; | ||
| protected HoodieSecondaryIndexInfo secondaryIndexMetadata; | ||
| protected Object vectorIndexMetadata; |
There was a problem hiding this comment.
🤖 nit: every other metadata payload field has its specific Avro type (HoodieMetadataBloomFilter, HoodieRecordIndexInfo, HoodieSecondaryIndexInfo, …), but this one is Object. Could you introduce a marker interface (e.g. HoodieVectorIndexEntry) that the Avro-generated types implement, or at least use a common base type? As-is, readers have no way to know what concrete types can land here without searching all constructors.
| 1, | ||
| randomSeed, | ||
| assumeNormalized, | ||
| partitionPath); |
There was a problem hiding this comment.
🤖 nit: this overload accepts quantizedCodeBytes, rabitqBits, and assumeNormalized but silently drops all three when delegating to the lower overload. Could the unused parameters be removed from the signature, or have a comment explaining why they're intentionally ignored? A caller passing a non-default rabitqBits here would have no indication their value is discarded.
|
|
||
| @Override | ||
| public String encode() { | ||
| return shardId == null |
There was a problem hiding this comment.
🤖 nit: the substring(0, 9) magic number is duplicated from HoodieTableMetadataUtil.getVectorIndexPostingPrefix — could this delegate to that helper, or at least reference a named constant for the cluster-prefix length? It's non-obvious why 9 bytes is the right cut-off here.
0682b18 to
6e11f97
Compare
|
Rebased this stack onto current Key correction:
Java 17 validation after the refresh:
The earlier internal v2/v3 terminology was intentionally removed: this is a new persisted format, not a migration from the experimental pre-shipping implementation. |
hudi-agent
left a comment
There was a problem hiding this comment.
Thanks for working on this! This PR adds an engine-neutral RaBitQ encoder and a metric-aware multibit scorer contract, along with the vector-index MDT key/routing plumbing and a supporting encoder/scorer certification test suite. The main thing worth confirming is the extended bit-plane weighting/ordering in RaBitQByteLutScorer and whether it stays consistent with the RaBitQ encoding across the posting-scan kernels. Please take a look at the inline comments, and this should be ready for a Hudi committer or PMC member to take it from here. Two small nits below — the rest of the code is clean and well-commented. a few notable readability concerns — primarily the silently-dropped parameters in a quantizer factory overload and the very high parameter counts in several factory/builder methods.
| } | ||
| double extendedDot = 0.0; | ||
| for (int p = 0; p < exBits; p++) { | ||
| extendedDot += (double) (1L << (exBits - 1 - p)) * planeDot(exBuffer, view.exPlaneOffset(vectorIndex, p)); |
There was a problem hiding this comment.
🤖 This weights extended plane p by 2^(exBits-1-p) (plane 0 = MSB), but RaBitQPlaneKernel.scorePass2 weights plane b by 2^b (plane 0 = LSB). Both docs claim to reconstruct the same centered code as multibitDotTerm. Could you confirm the posting-block exPlanes writer and the long[][] plane writer really emit opposite bit orderings? If a single repack ever feeds both with the same plane index, one path would silently return wrong scores.
| @Override | ||
| public String encode() { | ||
| return shardId == null | ||
| ? VectorIndexMetadataKey.postingPrefix(generationId, clusterId, 0).substring(0, 9) |
There was a problem hiding this comment.
🤖 nit: the literal 9 in .substring(0, 9) is opaque — could you extract a named constant (e.g. POSTING_PREFIX_SCAN_LENGTH) or add a short comment explaining what those 9 bytes cover (generation + cluster id, no shard)? A future developer changing the key layout will have no signal that this number needs updating.
| } | ||
|
|
||
| private static int[] exactTopK(float[] q, float[][] data, int k) { | ||
| float[] d = new float[data.length]; |
There was a problem hiding this comment.
🤖 nit: System.out.printf in a unit test gets lost or interleaved in CI logs and is inconsistent with how Hudi surfaces diagnostics — have you considered LOG.debug(...) or just dropping this line?
| return createVectorIndexQuantizerMetadataRecord( | ||
| quantizerType, | ||
| quantizedCodeBytes, | ||
| 1, |
There was a problem hiding this comment.
🤖 nit: the 6-param overload accepts quantizedCodeBytes and rabitqBits but then silently ignores both and delegates to the 3-param version — could you either use those arguments (pass rabitqBits instead of 1) or remove the overload entirely? A caller who passes a custom bits value will be surprised to find their setting discarded.
| return createVectorIndexQuantizerMetadataRecord(1, 0, quantizerType, randomSeed, null, partitionPath); | ||
| } | ||
|
|
||
| public static HoodieRecord<HoodieMetadataPayload> createVectorIndexQuantizerMetadataRecord( |
There was a problem hiding this comment.
🤖 nit: createVectorIndexManifestRecord has ~30 parameters, which makes call sites very hard to read and maintain. Have you considered introducing a VectorIndexManifestSpec value object (or a builder) to carry these fields? Even grouping the factor-config params (kappa, gMin, eps1Max, epsNRel) into the existing RaBitQFactorConfig would meaningfully reduce the arity here.
| rowPosition); | ||
| } | ||
|
|
||
| public PostingBlockBuilder addRow(String recordKey, |
There was a problem hiding this comment.
🤖 nit: the 14-parameter addRow overload is hard to call without a named-argument IDE. Have you considered introducing a small RowSpec (or builder) to bundle the scalar factors and location fields? Even just grouping the six scalar floats (fAdd1, fRescale1, err1, fAddEx, fRescaleEx, residualNorm) into an existing type like QuantizedVector would cut the arity significantly.
| String recordKey = VectorIndexMetadataKey.clusterStats(generation, clusterId); | ||
| HoodieVectorIndexClusterStats stats = new HoodieVectorIndexClusterStats( | ||
| routingVersion, | ||
| shardCount, |
There was a problem hiding this comment.
🤖 nit: fileGroupIds == null ? java.util.Collections.emptyList() : fileGroupIds.stream().collect(Collectors.toList()) — the fully-qualified java.util.Collections is unusual (there are imports already), and when fileGroupIds is non-null new ArrayList<>(fileGroupIds) is simpler and avoids the stream allocation.
Describe the issue this Pull Request addresses
Resolves #19098 (umbrella: #19094, RFC-109).
Third RFC-109 slice: an engine-neutral RaBitQ encoder and metric-aware, multibit scorer contract.
Summary and Changelog
RaBitQEncoder/QuantizedVector: deterministic 1-8 bit encoding, residual encoding, packed sign/extended codes, validated immutable encoded payloads, and finite-input checks.MetricQueryStateranking contracts for L2 (squared ranking distance), cosine, and dot product; validated query/rotation state; separateRaBitQDistanceScorerfor packed-code reconstruction.RaBitQByteLutScorer,RaBitQPlaneKernel, andVectorQueryPlanes, including persisted-plane ordering and posting-layout validation.VectorQuantizeroption enum; the old speculative runtime interface is replaced by concrete encoder/query/value types.Validation
Java 17 reactor validation:
Result: 82 tests, 0 failures/errors/skips; Checkstyle 0 violations; Apache RAT passed; residual multibit Recall@10 = 0.918.
Impact
Additive classes under
hudi-common; no write/read runtime path is wired in this PR. Bootstrap, lifecycle/publication, incremental maintenance, and query integration remain follow-up slices.Risk Level
Medium: isolated from existing runtime paths, but this establishes persisted-code and numeric contracts used by later slices.
Documentation Update
Contracts and multibit decomposition are documented in Javadoc. The user-facing RFC is tracked in #19309.
Contributor's checklist