(feat): PSDB Layout 2

[Adit2607]

What type of PR is this?

• 🍕 Feature
• 🐛 Bug Fix
• ⚙️ Config Change
• 📄 Template Change
• 📝 CHANGELOG/README Update
• 🎨 Linting
• 🧑‍💻 Code Refactor
• 🔥 Performance Improvements
• ✅ Test
• 🔁 Version Bump
• ⏩ Revert

---

📂 Modules Affected

• horizon (Real-time systems / networking)
• online-feature-store (Feature serving infra)
• trufflebox-ui (Admin panel / UI)
• infra (Docker, CI/CD, GCP/AWS setup)
• docs (Documentation updates)
• Other: ___________

---

Added tests

• 👍 Yes
• 🙅 No
• 🙋 Need Help!

• 👍 Manual
• 📝 Unit Tests
• 🔗 API/Integration Tests

Key Requirement Doc

Downstream Impact

• No downstream impact

Additional Information

Updated by ReadyToReviewBot for @shubhamk-meesho · Edit form

Summary by CodeRabbit

• New Features

  • Introduced Layout V2 storage format with bitmap-based optimization for sparse feature data, enabling improved compression and storage efficiency.
  • Added layout comparison testing infrastructure to measure performance and efficiency gains across storage formats.

• Configuration

  • New shadow comparison configuration to enable performance testing and monitoring between storage layouts.

[coderabbitai]

Walkthrough

This pull request introduces a Layout V2 serialization format for Permanent Storage Data Blocks with bitmap-based sparse feature encoding, abstract PSDBBlock interface to unify Layout V1/V2 deserialization, and shadow comparison infrastructure to evaluate Layout 2 benefits. API updates thread bitmap data through feature parsing and retrieval chains, while configuration extensions enable layout comparison sampling.

Changes

Cohort / File(s)	Summary
PSDBBlock Interface & Core Abstraction online-feature-store/internal/data/blocks/deserialized_psdb.go	New PSDBBlock interface defining contract for deserialized blocks with accessor methods (GetDataType, GetLayoutVersion, IsExpired, etc.) and feature extraction methods across all data types/layouts.
Layout V1 Implementation online-feature-store/internal/data/blocks/deserialized_psdb_v2.go, deserialized_psdb_v2_test.go	Updated Layout V1 PSDB deserialization to return PSDBBlock interface; added accessor methods; refactored feature getters to include defaultValue parameters; replaced Copy() with CopyBlock(); updated test assertions to use interface getters.
Layout V2 Serialization online-feature-store/internal/data/blocks/perm_storage_datablock_v3.go	New Layout V2 serialization with 17 type-specific serializers; supports bitmap-aware sparse encoding, per-element packing, vector handling, boolean bit-packing, and bitmap prepending logic.
Layout V2 Deserialization online-feature-store/internal/data/blocks/deserialized_psdb_v3.go	New DeserializedPSDBLayout2 type with bitmap-aware feature extraction; conditionally uses bitmap sparse access when present, falling back to V1 dense logic; includes bitmap validation and dense offset computation helpers.
Serialization Constants & Builder online-feature-store/internal/data/blocks/perm_storage_datablock_v2.go, psdb_builder.go, psdb_pool_test.go	Renamed PSDBLayout1LengthBytes to PSDBLayout1HeaderBytes (9 bytes); added PSDBLayout2HeaderBytes (10 bytes) and bitmap bit flags; extended PermStorageDataBlock with bitmap field and SetBitmap/SetupBitmapMeta builder APIs; updated all buffer sizing logic.
Shadow Layout Comparison online-feature-store/internal/data/blocks/psdb_shadow_compare.go, layout_comparison_test.go, layout_comparison_results.txt	Added ShadowSerializeAsLayout function for off-path Layout 1 vs Layout 2 size comparison; comprehensive test framework with TestLayout1VsLayout2Compression and TestLayout2BitmapOptimization; generated ASCII results report with per-type metrics and recommendations.
Cache Storage online-feature-store/internal/data/blocks/cache_storage_datablock_v2.go, cache_storage_datablock_v2_test.go	Updated FGIdToDDB map from *DeserializedPSDB to PSDBBlock interface; refactored serialization handlers for Layout 1/2; replaced direct field access with accessor methods; integrated NegativeCacheExpiredPSDB helper.
Serialization Tests online-feature-store/internal/data/blocks/perm_storage_datablock_v2_test.go, perm_storage_datablock_v2_bench_test.go	Updated all buffer boundary constants from PSDBLayout1LengthBytes to PSDBLayout1HeaderBytes across deserialization and serialization test cases.
Feature Data Parsing online-feature-store/internal/system/system.go	Extended ParseFeatureValue and all Get* methods to return (data, bitmap \[\]byte, error); bitmap signals non-default values per feature; added per-vector default tracking, comparison helpers, and metadata lookups by sequence.
Persistence Handler online-feature-store/internal/handler/feature/persist.go, persist_test.go	Updated BuildPSDBBlock to accept featureBitmap parameter and set it via SetBitmap; conditionally calls SetupBitmapMeta for Layout 2; added maybeShadowCompare helper to emit size reduction metrics based on config sample rate.
Retrieval Handler online-feature-store/internal/handler/feature/retrieve.go, model.go	Changed FGData and RetrieveData maps from *DeserializedPSDB to PSDBBlock interface; updated fillMatrix and feature retrieval to use accessor methods and thread defaultValue through all paths; replaced field mutations with SetNegativeCache setter.
Data Repository Stores online-feature-store/internal/data/repositories/stores/redis.go, scylla.go, store.go	Updated RetrieveV2 and BatchRetrieveV2 to return PSDBBlock interface instead of pointer; refactored internal maps and deserialization logic to use interface accessors and NegativeCacheDeserializePSDB/NegativeCacheExpiredPSDB helpers.
Configuration Extensions online-feature-store/internal/config/config.go, etcd.go, mock.go, models.go	Added LayoutShadowComparisonConfig struct with Enabled and SampleRate fields; added GetLayoutShadowComparisonConfig() method to Manager interface and implementations; added LayoutShadowComparison field to FeatureRegistry.
Utilities .gitignore	Added .claude/ directory to ignore patterns.

🚥 Pre-merge checks | ✅ 1

✅ Passed checks (1 passed)

Check name	Status	Explanation
Dynamic Configuration Validation	✅ Passed	No dynamic configuration files (application-dyn-*.yml) exist in the repository; PR modifies Go config structs only.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

♻️ Duplicate comments (2)

online-feature-store/internal/data/blocks/layout_comparison_results.txt (1)

1443-1443: ⚠️ Potential issue | 🟡 Minor

Fix the double percent sign typo.

Line 1443 contains 0%% which should be 0%.

📝 Proposed fix

-  • Overhead only in edge cases with 0%% defaults
+  • Overhead only in edge cases with 0% defaults

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@online-feature-store/internal/data/blocks/layout_comparison_results.txt` at
line 1443, Replace the stray double-percent typo "0%%" in the text snippet "•
Overhead only in edge cases with 0%% defaults" with a single percent sign so it
reads "0%"; search for the literal "0%%" occurrence and update it to "0%" (no
other changes).

online-feature-store/internal/data/blocks/deserialized_psdb_v2.go (1)

339-354: ⚠️ Potential issue | 🟡 Minor

Missing bounds check for pos parameter persists.

Per the past review comment, GetNumericVectorFeature still lacks an explicit bounds check for pos < 0 || pos >= len(vectorLengths) before accessing vectorLengths[pos] at line 349. A negative pos would cause the loop to never match and then panic.

This was acknowledged but not addressed in this PR. Consider adding the defensive check:

🛡️ Proposed fix

 func (dd *DeserializedPSDB) GetNumericVectorFeature(pos int, vectorLengths []uint16, defaultValue []byte) ([]byte, error) {
+	if pos < 0 || pos >= len(vectorLengths) {
+		return nil, fmt.Errorf("position %d out of bounds for vectorLengths (len=%d)", pos, len(vectorLengths))
+	}
 	data := dd.OriginalData
 	size := dd.DataType.Size()
 	var start int = 0

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@online-feature-store/internal/data/blocks/deserialized_psdb_v2.go` around
lines 339 - 354, GetNumericVectorFeature is missing a defensive bounds check for
the pos parameter; add an explicit check at the start of the function to verify
pos >= 0 and pos < len(vectorLengths) and return a clear error (e.g. "position
out of bounds") if it fails, before the loop or any access to
vectorLengths[pos], so the later loop and the computation of end cannot panic
when pos is negative or >= len(vectorLengths).

🧹 Nitpick comments (2)

online-feature-store/internal/data/blocks/perm_storage_datablock_v2.go (1)

59-86: Potential issue: Clear() reads layoutVersion before resetting it.

The function captures headerLen based on p.layoutVersion (lines 60-63), then resets p.layoutVersion = 0 (line 64). This is correct for determining the current buffer trim size. However, if Clear() is called on a freshly allocated struct where layoutVersion == 0, headerLen defaults to PSDBLayout1HeaderBytes (9), which may not match the actual buffer size if it was pre-allocated for layout 2.

This is likely fine since the builder initializes the buffer appropriately, but worth noting for defensive coding.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@online-feature-store/internal/data/blocks/perm_storage_datablock_v2.go`
around lines 59 - 86, PermStorageDataBlock.Clear determines headerLen from
p.layoutVersion which can be zero for newly allocated structs, causing an
incorrect trim; update Clear (function PermStorageDataBlock.Clear) to infer
headerLen defensively from the buffer instead of only from p.layoutVersion: if
p.layoutVersion != 0 use it as now, otherwise inspect len(p.buf) and choose
PSDBLayout2HeaderBytes when buf is large enough, fall back to
PSDBLayout1HeaderBytes or to len(p.buf) if smaller; then proceed to reset
p.layoutVersion and trim p.buf accordingly. Ensure references to
PSDBLayout1HeaderBytes, PSDBLayout2HeaderBytes, p.layoutVersion and p.buf are
used in the new logic.

online-feature-store/internal/handler/feature/persist.go (1)

6-6: Consider using math/rand/v2 or explicit seeding for deterministic behavior in tests.

math/rand uses a global source that defaults to a deterministic seed (1) in Go <1.20 and auto-seeds in Go ≥1.20. For production sampling this is fine, but be aware that test determinism may vary across Go versions.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@online-feature-store/internal/handler/feature/persist.go` at line 6, The
import of "math/rand" should be made explicit for deterministic tests: either
switch the import to "math/rand/v2" and update any calls, or (preferably) make
randomness injectable or explicitly seed the global source; locate uses of rand
(e.g., rand.Intn, rand.Float64) in persist.go and then either (A) replace global
calls with an injected *rand.Rand instance (created with
rand.New(rand.NewSource(seed)) in tests) or (B) add an init() that calls
rand.Seed(time.Now().UnixNano()) for production and ensure tests call
rand.Seed(fixedSeed) or use rand.NewSource(fixedSeed) to guarantee deterministic
behavior across Go versions.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Duplicate comments:
In `@online-feature-store/internal/data/blocks/deserialized_psdb_v2.go`:
- Around line 339-354: GetNumericVectorFeature is missing a defensive bounds
check for the pos parameter; add an explicit check at the start of the function
to verify pos >= 0 and pos < len(vectorLengths) and return a clear error (e.g.
"position out of bounds") if it fails, before the loop or any access to
vectorLengths[pos], so the later loop and the computation of end cannot panic
when pos is negative or >= len(vectorLengths).

In `@online-feature-store/internal/data/blocks/layout_comparison_results.txt`:
- Line 1443: Replace the stray double-percent typo "0%%" in the text snippet "•
Overhead only in edge cases with 0%% defaults" with a single percent sign so it
reads "0%"; search for the literal "0%%" occurrence and update it to "0%" (no
other changes).

---

Nitpick comments:
In `@online-feature-store/internal/data/blocks/perm_storage_datablock_v2.go`:
- Around line 59-86: PermStorageDataBlock.Clear determines headerLen from
p.layoutVersion which can be zero for newly allocated structs, causing an
incorrect trim; update Clear (function PermStorageDataBlock.Clear) to infer
headerLen defensively from the buffer instead of only from p.layoutVersion: if
p.layoutVersion != 0 use it as now, otherwise inspect len(p.buf) and choose
PSDBLayout2HeaderBytes when buf is large enough, fall back to
PSDBLayout1HeaderBytes or to len(p.buf) if smaller; then proceed to reset
p.layoutVersion and trim p.buf accordingly. Ensure references to
PSDBLayout1HeaderBytes, PSDBLayout2HeaderBytes, p.layoutVersion and p.buf are
used in the new logic.

In `@online-feature-store/internal/handler/feature/persist.go`:
- Line 6: The import of "math/rand" should be made explicit for deterministic
tests: either switch the import to "math/rand/v2" and update any calls, or
(preferably) make randomness injectable or explicitly seed the global source;
locate uses of rand (e.g., rand.Intn, rand.Float64) in persist.go and then
either (A) replace global calls with an injected *rand.Rand instance (created
with rand.New(rand.NewSource(seed)) in tests) or (B) add an init() that calls
rand.Seed(time.Now().UnixNano()) for production and ensure tests call
rand.Seed(fixedSeed) or use rand.NewSource(fixedSeed) to guarantee deterministic
behavior across Go versions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 5ebe8c3d-6d77-4bbe-8d02-62c692be9364

📥 Commits

Reviewing files that changed from the base of the PR and between 72315d5 and 7437301.

📒 Files selected for processing (7)

• online-feature-store/internal/data/blocks/cache_storage_datablock_v2.go
• online-feature-store/internal/data/blocks/deserialized_psdb_v2.go
• online-feature-store/internal/data/blocks/deserialized_psdb_v3.go
• online-feature-store/internal/data/blocks/layout_comparison_results.txt
• online-feature-store/internal/data/blocks/perm_storage_datablock_v2.go
• online-feature-store/internal/data/blocks/psdb_shadow_compare.go
• online-feature-store/internal/handler/feature/persist.go

✅ Files skipped from review due to trivial changes (1)

• online-feature-store/internal/data/blocks/deserialized_psdb_v3.go

🚧 Files skipped from review as they are similar to previous changes (2)

• online-feature-store/internal/data/blocks/psdb_shadow_compare.go
• online-feature-store/internal/data/blocks/cache_storage_datablock_v2.go