Our Feature Store Saga: From Spark to Scale (part 4)

Hi, I am Daniil, Staff Data Engineer from Data Services Team and this is Pimp My ML Process Show.

In the previous chapters of our story, we already described the first stable version of our DWH, challenge of data sync and data exchange crunching.

1. The Spark: An Idea Ignites 🔥

Every scaling journey has its friction points. Ours was in ML pipelines. Teams were duplicating effort, building similar (or worse, slightly different) features on top of our DWH (BigQuery). This wasn’t just inefficient; it introduced significant risk and slowed us down.

The problem wasn’t just computing on the DWH; it was that while training datasets were collected there, on-demand models couldn’t access that data in production. Consequently, teams had to rewrite SQL to directly go to microservices’ databases because their APIs weren’t evolving fast enough. If an upstream service changed a table schema, the ML service consuming it could instantly break.

We critically faced training-serving skew, a major form of data drift. Models would perform well in “teaching” (training) but then fail or behave unpredictably in production (inference). We desperately needed a single source of truth to sync the data and logic between our training and inference environments.

The core needs that sparked the Feature Store idea:

• Consistency: Uniform definitions and calculations. Most importantly, we needed to guarantee that the exact same feature logic was used for both model training and serving to eliminate this data drift.
• Discoverability: A clear catalog of existing features.
• Efficiency: Build once, reuse by everyone.

While off-the-shelf solutions existed, a custom build offered strategic control over our data, cost-efficiency, and alignment with our specific tech landscape. The vision: a baseline service, not just a tool that is optionally used.

2. v0.1: Laying the First Stones 🛠️

Our MVP mantra: “simple and functional.” We focused on core essentials.

Foundational Requirements:

• Basic Feature Registry: A simple catalog is a basic DB schema.
• Offline (Cold) Features: Scheduled calculations (daily/hourly) for batch needs.
• Basic Feature Retrieval: Instructions for accessing feature tables.
• Manual Feature Addition Process: Accessing database and adding features.

Initial Tech Stack:

• Google BigQuery: Primary data source and feature storage.
• Airflow: Orchestrator for batch feature calculations.
• PostgreSQL: The serving layer where features landed after calculation.

Even this basic version brought value:

• Reduced some feature redundancy.
• Forced clarity in feature definitions.
• Showed the potential for a more streamlined MLOps process.

It was awkward to use, with manual adds and retrievals, relied on BigQuery access and solely on batch calculations, but it was a start.

Feature store architecture v1

3. The Strain: When Success Revealed Cracks 🧱

Increased usage quickly highlighted the MVP’s limitations. The initial warning signs grew louder.

Performance & Scalability Hits:

• Calculation Latency: Not yet established SLAs were already not satisfied.
• Serving Latency: Direct BigQuery access led to contention and slow queries.

Outpaced by Data Science Needs:

• Growing demand for streaming and on-demand feature calculations.
• Manual and error-prone dataset creation.
• Feature addition: the review and integration process became a bottleneck. It wasn’t scalable.
• Unified API serving endpoint with flexible and fast serving.

Moreover we weren’t just dealing with generic ML models. We faced a dichotomy of needs. Credit Scoring, Recommendation Engines, AFS (Anti-Fraud Systems) — all imposed different requirements on consistency, real-time calculations, flexibility of calculations. One size definitely did not fit all.

Our backlog rapidly filled. The MVP had served its purpose; an evolution was critical.

4. The Rebuild: Architecting for Tomorrow ⚙️

We needed a fundamental shift to handle diverse feature types, speed, and scale. Our new design version emerged, focusing on modularity.

Key Architectural Upgrades:

• Embracing the Stream with Benthos: Introduced for resilient, high-throughput “Stream Feature” processing from Pub/Sub, achieving sub-minute freshness. The demand for features reacting to events from Pub/Sub topics was high. Airflow batch jobs weren’t sufficient for this. After researching streaming processing tools which is capable and GCP-friendly we brought in Benthos, a fantastic stream processing tool. This allowed us to build resilient, high-throughput pipelines that could consume from Pub/Sub, perform transformations, and land these fresh features into stores ready for quick retrieval.
• Live clicks with ClickHouse: For “Live Features” — those computed on-demand using real-time data — ClickHouse, when properly tuned and with well-designed queries, offered incredible speed. We defined clear patterns for its use, focusing its power where it performed best: aggregating fresh data from its own fast storage for those sub-second latency requirements. Now it was possible to calculate ad-hoc queries as features from the same api.
• A multi-database approach — The Right Tool for the Job:

• Matured Component Ecosystem: Airflow (batch), Pub/Sub (events), Benthos (streaming), ClickHouse (live), Postgres (registry and serving), Bigtable (fast serving experiment), and dedicated services (GCP Cloud Run/GKE) for APIs and compute logic.

This overhaul allowed independent scaling of parts and improved overall system reliability.

Feature store architecture v2

5. At full capacity: The Feature Store At Date ✨

The platform is now a central, capable part of our MLOps.

Current Strengths:

• Comprehensive Feature Support: Robust handling of Cold, Stream, and Live features.
• Sophisticated Registry: PostgreSQL-backed with rich metadata.
• Streamlined Dataset Creation: Tools and SDKs for easier training data assembly.
• Standardized Retrieval: A documented “Features Store API” for programmatic access.
• Improved Feature Addition: More formalized process via YAML & CI/CD.
• Meeting SLAs: Established SLI/SLAs. Generally hitting targets for calculation/serving latency, availability, and pipeline reliability. Each type has relevant targets.
• Enhanced User Experience: The UI for users and admins, expressive documentation offer significantly better usability.

6. The Horizon: What’s Next on the Map? 🗺️

We’re not standing still. Key future initiatives:

• Relentless Cost Optimization: This is a big one. Running a sophisticated, multi-layered data platform means keeping a close eye on the bills. We’ll be diving deeper into optimizing our storage strategies — are we using the most cost-effective tiers in BigQuery, ClickHouse, and Bigtable for different data access patterns? Can we be smarter about our compute usage in Airflow, Benthos, and our serving services? Every query, every pipeline, every stored byte will be carefully reviewed.
• Supercharged Config Generation & Management: While adding features is way better than it used to be, we think we can make it even smoother. We’re looking at more advanced config files generation tools and processes. Imagine a world where defining a new feature, its transformations, its storage, and its monitoring is almost entirely declarative, with intelligent defaults and robust validation catching errors before they even hit dev. This ties into making the whole system easier to manage and less prone to manual mistakes.
• Enhanced UI/UX: Continuously improving the UI to make feature discovery, understanding lineage, and even contributing new features more intuitive. Clearer Guardrails: As we open things up, we need robust guardrails — better automated checks for data quality and freshness, performance implications of new features, and adherence to governance standards.
• Advanced Monitoring & Alerting: We’re pretty good here, but we want to get great. Think more proactive alerting on potential data quality issues before they impact models, and deeper insights into feature usage and drift.
• Evolving Governance & Roles: As more teams contribute, refining our Responsibilities matrix and governance model will be key to maintaining order and quality.
• Make building Feature Store like a Lego: We aim to extract core components and make all parts of the system “pluggable”. So other teams in the company could start their own version of FS from ready bricks.
• Optimisation of calculations: live and streaming feature types have become very popular. it increased demand for more throughput and larger aggregate window sizes. This is a tough point on the edge of what is possible to achieve using current level of technology.
• Exploring New Tech & Techniques: We’re always curious. Are there new processing engines, storage solutions, or MLOps techniques (like more sophisticated feature encoding or an even tighter integration with our model training/serving platforms) that could give us an advantage? For example, that “Calculating metric in a hierarchical manner” idea is something we’re eager to explore for certain complex feature types. The goal isn’t just to add more features to the Feature Store, but to make the entire lifecycle of a feature — from idea to production value — faster, more reliable, and more collaborative. It’s an ongoing journey, but we’re excited about where we’re headed.

Our aim: make the entire feature lifecycle faster, more reliable, and more collaborative.

7. The Debrief: Wisdom from the Front Lines 🎓

Our journey from a simple idea to a core platform taught us several key lessons:

• Start Simple, Design for Growth: MVPs are vital, but early architectural foresight for scale can prevent significant rework.
• A Feature Store is a Product: It requires continuous iteration, user-focus, and a dedicated roadmap.
• Abstractions Over Raw Tech: The specific tools are important, but user-friendly abstractions (APIs, SDKs, UI) are what deliver value from underlying complexity.
• People & Process are Paramount: Adoption, clear roles, responsibilities matrix, and streamlined processes for development, inter team communication, incidents tactics are as critical as the technology.

Building our Feature Store was a significant undertaking, but its impact on our ML velocity, reliability, and innovation has made it undeniably worth it. It’s a foundation for future intelligence, and the evolution continues.