ML Feature Pipelines and Feature Stores Questions

Direct answer

Immediate mitigations buy time without a code change: shorten the watermark threshold (drops old, rarely-needed state faster) and, if the streaming query supports it, force an early checkpoint/restart with a tighter configuration to release currently-held state. Long-term, the real fix is almost always adding or shortening a TTL (time-to-live) on state that should not live forever, since unbounded state growth is usually a symptom of state accumulating for keys that will never see another matching event, not a sizing problem alone; if the state genuinely needs to outlive what Spark's in-memory/RocksDB state store can hold efficiently, moving to an EXTERNAL state store (or redesigning the windowing to not need unbounded per-key history at all) is the architectural escape valve.

Structured elaboration

Why state grows unbounded in the first place. Spark's Structured Streaming state store holds accumulating per-key state for stateful operations (a windowed aggregation's running totals, a flatMapGroupsWithState custom aggregation's accumulated values, a stream-stream join's buffered unmatched rows) across micro-batches; WITHOUT a watermark (or with one set too generously wide), state for a key is never considered "safe to drop," so keys that appear once and never again still hold their state entry FOREVER, and the state store's total footprint grows monotonically with the number of DISTINCT keys ever seen, not the number of ACTIVE keys at any given time.

Immediate mitigations.

• Shorten the watermark threshold. A watermark that is wider than the REAL lateness distribution needs (sizing it against real data) holds state open longer than necessary; tightening it to match the ACTUAL observed lateness (not an arbitrarily generous default) directly bounds how long a window's state stays open before being released.
• State compaction / a forced checkpoint cycle. RocksDB-backed state stores (the default for larger stateful workloads) accumulate internal fragmentation over many micro-batches; triggering compaction (spark.sql.streaming.stateStore.rocksdb.compactOnCommit or a similar maintenance operation, depending on Spark version) reclaims space from already-tombstoned (deleted-but-not-yet-reclaimed) state entries without needing a full restart.

Long-term strategies.

• TTL eviction. For stateful operations expressed via flatMapGroupsWithState/applyInPandasWithState (arbitrary custom stateful logic, this API), explicitly setting a TIMEOUT (event-time or processing-time based) on each group's state ensures a key with no new events for the configured duration has its state EXPLICITLY cleared, rather than relying solely on watermark-driven window closure (which only applies to the built-in windowed-aggregation shape, not arbitrary custom state).
• Incremental snapshots. Rather than a full state-store snapshot on every checkpoint (expensive and slow for very large state), Spark's state store supports INCREMENTAL checkpointing (only the DELTA since the last snapshot), reducing both the per-checkpoint cost and the RECOVERY-time cost of replaying from a checkpoint, a meaningful operational win once state size grows large enough that full snapshots themselves become a bottleneck.
• Windowing changes. If the query's CURRENT design needs per-key state to live indefinitely because the business logic is framed as "track this forever," reframing to a BOUNDED window (a rolling 30-day window instead of "all time," for instance) if the business requirement genuinely tolerates it directly bounds state size by construction, the most durable fix when it is available.
• External state store. If the required state genuinely cannot be bounded by any of the above (a business requirement to track something truly indefinitely, at a scale beyond what an in-cluster state store can hold efficiently), moving that state to an EXTERNAL, purpose-built store (a key-value store queried by the stateful operation instead of held in Spark's own state store) is the architectural escape valve, at the cost of added system complexity (network round-trips to the external store per micro-batch, a new dependency to operate) that should be reached for only once the simpler, in-Spark options are genuinely insufficient.

Worked example

A Structured Streaming job tracking "days since a user's last purchase" via flatMapGroupsWithState, keyed by user_id, with NO timeout configured: every user who EVER made a purchase, even years ago and never again, still holds an open state entry, and the state store's total size grows with the CUMULATIVE user count over the query's entire lifetime, not the count of users who are still ACTIVE.

Immediate mitigation: if a watermark-based windowed aggregation were ALSO present in the same query, tightening its threshold helps that PART of the state, but does nothing for the flatMapGroupsWithState state specifically, since watermark-driven closure only applies to the built-in windowed-aggregation shape.

Long-term fix: add an explicit event-time timeout to the flatMapGroupsWithState call (e.g., GroupStateTimeout.EventTimeTimeout with a 180-day threshold): a user with no purchase activity for 180 days has their state entry explicitly cleared on the next batch that crosses that threshold, bounding the state store's size to roughly "distinct users active within the last 180 days," not "every user ever seen," directly addressing the actual root cause (a missing TTL on genuinely stateful, indefinitely-accumulating tracking) rather than just buying time with tighter checkpointing.

Trade-offs and pitfalls

• Common mistake: tightening the watermark and declaring the problem solved, when the actual unbounded growth is coming from a flatMapGroupsWithState operation the watermark does not govern at all; diagnosing WHICH stateful operation is actually accumulating (via the Spark UI's Structured Streaming state-store metrics, which report size PER operator) before choosing a fix avoids applying the wrong lever.
• Common mistake: reaching for an external state store as the first response to state growth, when a properly-configured TTL/timeout on the existing state operations would have solved it with far less added complexity; treat the external store as the LAST option once the simpler, in-Spark fixes are confirmed insufficient, not the first instinct.
• A TTL that is too aggressive drops state for a genuinely still-relevant key too early, a correctness cost (losing legitimate history) traded for a memory savings; size the TTL against the REAL business meaning of "this key is no longer relevant," not purely against memory pressure.
• Incremental snapshots reduce checkpoint COST but do not by themselves bound state SIZE; they help the operational overhead of an already-large state store, but do not address the underlying unbounded-growth root cause the way a TTL/timeout does, worth applying alongside a real bounding fix, not as a substitute for one.

Direct answer: Hitting a 5ms median at 100k requests/sec means the design has to eliminate every unnecessary network hop and put the hot path entirely in memory close to the caller, with the data store itself sharded wide enough that no single shard is the bottleneck. The architecture layers, in order from the client: a local or co-located cache, a sharded low-latency store (Redis Cluster or a DynamoDB-style store), multi-region replicas, and an observability layer that separately tracks median and tail latency because at this scale they diverge.

Structured elaboration:

flowchart LR
  Client["Inference client"] --> Region1["Regional read path"]
  subgraph Region["Nearest region"]
    Region1 --> LocalCache["Local cache"]
    LocalCache -->|miss| Shard["Sharded online store (20 shards, 3x replicas)"]
  end
  Primary["Primary write path"] -->|async replication| Region
  Primary -->|async replication| RegionB["Other regions"]
  Shard -->|slow path fallback| Stale["Serve slightly stale cached value"]

1. Data store and sharding. Partition the keyspace (typically by hashing entity ID) across enough shards that each shard handles a small fraction of 100k QPS (queries per second) with headroom, e.g. 20 shards at ~5k QPS each with 3x replication for read fan-out.
2. Multi-region caching. Serve reads from the region closest to the caller. Each region holds a local read replica (or a fully local cache warmed from the primary) so a request never crosses a region boundary on the read path; cross-region replication happens asynchronously in the background.
3. Cache warming. On a region failover or a cold cache, warming from a snapshot or from the primary region avoids a thundering herd of cache misses hitting the backing store simultaneously; a common pattern is to replay the last N minutes of writes into the new cache before routing traffic to it.
4. Hotspot mitigation. A small number of keys will still dominate traffic (see the hot-key sub-area); mitigate with a local in-process cache for the hottest keys, request coalescing (collapse concurrent identical lookups into one backend call), and, if needed, replicating hot keys onto every shard rather than one.
5. Consistency model. Accept eventual consistency between regions (bounded by replication lag, typically tens to low hundreds of milliseconds) in exchange for the latency win; document the staleness bound so consumers know what they are trading.
6. Load balancing. Client-side or sidecar-based load balancing with health checks and fast failure detection, so a single slow shard does not drag down the aggregate tail.
7. Measuring and maintaining tail latency. Track p50, p95, p99, and p99.9 separately, not just an average; instrument per-shard and per-region so you can find the shard causing the tail rather than only seeing the aggregate. Use load shedding or a fast timeout-plus-fallback (serve a slightly stale cached value) for the small percentage of requests that would otherwise blow the SLA (service-level agreement).

Worked example: At 100k QPS with a target median of 5ms, budget the latency: roughly 1-2ms for the network hop to a co-located cache, 0.5-1ms for the cache lookup itself, and the remainder as safety margin for queueing under load. If each shard node can sustain 8,000 ops/sec at that latency (a realistic number for an in-memory store with pipelining), you need at least ceil(100,000 / 8,000) = 13 shards for the primary write path, and you would typically round up to 16-20 to leave headroom for uneven key distribution and node maintenance.

Trade-offs & pitfalls: The most common mistake at this scale is designing for the median and being surprised by the tail: a 5ms median with a 200ms p99.9 still means thousands of slow requests per second at 100k QPS, which is often unacceptable for a synchronous inference path. A second common mistake is under-provisioning replica read fan-out, so a single popular shard saturates even though the aggregate cluster has spare capacity. Cross-region async replication means a model can read a feature that is a few hundred milliseconds stale right after a write; if the use case cannot tolerate that (for example, a fraud rule keyed on a just-written flag), the design needs a synchronous or quorum-read path for that specific feature instead of the default async replication.

Direct answer: Hot-key mitigation for online feature serving combines three layers that compose: caching close to the caller to absorb most of the skewed traffic, replicating the hot keys themselves so no single shard owns the whole load, and routing/throttling changes so the system degrades gracefully instead of falling over when a key gets even hotter than expected.

Structured elaboration:

1. Caching. Add a local (in-process or co-located) cache in front of the online store for the hottest keys. Because the traffic is so skewed (top 0.1% of keys getting 90% of reads), even a small cache with a short TTL (time-to-live) captures most of the benefit: a cache holding only the top few thousand keys can absorb the majority of read volume.
2. Replication of hot keys specifically. Rather than replicating every key uniformly, detect the hot ones (see detection below) and replicate just those onto multiple shards or nodes, so reads for a hot key fan out across several replicas instead of hammering one.
3. Request routing. Route requests for a hot key round-robin or hash-based across its replicas, and route requests for cold keys normally; this requires the routing layer to be hot-key-aware rather than using a single static hash.
4. Adaptive prefetching. For predictable hot keys (a trending item, a celebrity user), prefetch and pin them into cache ahead of the traffic spike rather than reacting after the first cache miss storm.
5. Throttling and backpressure. As a last line of defense, apply per-key or per-shard rate limiting so one runaway hot key cannot starve the rest of the traffic on the same shard; combine with a fallback response (serve a slightly stale cached value or a default) rather than failing the request outright.
6. Detection. Continuously track per-key or per-shard request counts (a sliding-window counter or a sketch like Count-Min Sketch for cheap approximate top-K tracking) so the system can flag a new hot key automatically rather than relying on someone noticing degraded latency first.

Worked example: If the top 0.1% of user IDs account for 90% of reads at 100k QPS (queries per second), that is roughly 90,000 QPS concentrated on a key set that might be only a few thousand entities. A cache holding those few thousand keys, even with a modest hit rate of 95%, removes on the order of 85,000 QPS from the backing store, turning a saturating hot-partition problem into a manageable ~15,000 QPS of cache-miss traffic spread across the full keyspace.

Trade-offs & pitfalls: Caching hot keys introduces a staleness window; for features where freshness matters (a just-updated fraud flag), a short TTL or a write-through cache invalidation is needed instead of a pure time-based expiry, which adds complexity. Replicating hot keys onto multiple shards breaks the simple "one key, one owner" mental model and complicates writes: a write to a hot key now has to fan out to every replica, or you accept eventual consistency across them. Detection systems based on a fixed threshold can be slow to react to a sudden spike (a key that goes from cold to hot in seconds, e.g. breaking news); a sketch-based or exponentially-weighted detector reacts faster than a simple rolling count but adds its own tuning surface.

Direct answer: An aggressive watermark (a short allowed-lateness) closes windows and emits results quickly, minimizing latency but dropping or mishandling more genuinely late-arriving events; a conservative watermark (a longer allowed-lateness) waits longer before closing a window, capturing more late data at the cost of higher latency and more state held in memory.

Structured elaboration: The watermark is max_observed_event_time - allowed_lateness. Setting allowed-lateness to, say, 5 seconds means the window closes and emits almost immediately once the wall clock catches up, which is ideal for a low-latency online feature but means any event delayed by more than 5 seconds is either dropped or must be handled as a separate late-correction path. Setting it to 10 minutes means the window waits substantially longer before finalizing, capturing the vast majority of realistically-delayed events, but the feature is correspondingly less fresh, and the engine has to keep 10 minutes of window state open rather than a few seconds.

Worked example in Spark Structured Streaming: with .withWatermark("event_time", "10 minutes") on a 1-minute tumbling window, a window covering 10:00-10:01 will not emit its final result until the watermark passes 10:01, which happens once the max observed event time reaches roughly 10:11 (10 minutes of allowed lateness past the window end); an event for that window arriving after the watermark has already passed it is dropped from that aggregate (Structured Streaming's default behavior) unless the job is also configured to emit and merge late updates. Contrast that with a 10-second watermark on the same window: it closes and emits by roughly 10:01:10, over 100x faster, but discards essentially any event more than 10 seconds late.

Trade-offs & pitfalls: The biggest pitfall is picking a single watermark setting for a pipeline without measuring the actual lateness distribution of the source data; a watermark set shorter than the p99 lateness of real traffic silently and systematically drops a meaningful fraction of late data every window, which is easy to miss unless the pipeline monitors dropped-late-event counts explicitly. Conversely, an overly conservative watermark on a low-latency use case (fraud scoring needing sub-second freshness) can make the feature too stale to be useful even though it is more "complete," so the right choice depends on which side of the latency/completeness trade-off the specific feature's consumer actually needs, and different features in the same pipeline may legitimately need different watermark settings.

Direct answer: At-most-once means an event is processed zero or one times (never redelivered, so failures cause silent data loss); at-least-once means an event is processed one or more times (failures trigger redelivery, so duplicates are possible); exactly-once means the effect of processing is as if each event were applied precisely once, even though the underlying delivery mechanism may redeliver.

Structured elaboration:

• At-most-once: a producer sends an event and does not wait for or retry on failure (fire-and-forget). Practical example: a Kafka producer configured with acks=0. Achieved simply by not implementing retries; the trade-off is accepted data loss on any transient failure, which is rarely acceptable for feature pipelines feeding models.
• At-least-once: the producer retries until it gets an acknowledgment, and the consumer commits its offset only after successfully processing an event, so a crash between processing and committing causes that event to be reprocessed. Practical example: a Kafka consumer using manual offset commits with acks=all on the producer side. This is the most common default in Spark Structured Streaming and Flink without additional guarantees layered on.
• Exactly-once: achieved not by preventing redelivery (which is generally not fully preventable in a distributed system) but by making the consumer's processing idempotent, so redelivering the same event has no additional effect. Two common implementation strategies: (1) idempotent writes keyed by a stable event ID (an upsert that overwrites with the same result regardless of how many times it runs), or (2) transactional sinks that commit the output and the offset together atomically (Kafka's transactional producer/consumer API, or Flink's two-phase-commit sink), so either both the write and the offset advance, or neither does.

Worked example: A Spark Structured Streaming job writing feature aggregates to a database achieves effectively-exactly-once semantics not through Spark's delivery guarantee alone (which is at-least-once on the read side by default) but by writing with an idempotent upsert keyed by (window, entity_id), so if the same micro-batch is reprocessed after a failure, the second write produces the identical row rather than double-counting; this is the standard pattern rather than relying on a fully transactional end-to-end pipeline, which is harder to achieve across heterogeneous systems (Kafka to Spark to an external store).

Trade-offs & pitfalls: "Exactly-once" is a commonly overclaimed term; most systems that advertise it actually provide exactly-once effect through idempotency or transactions, not literally exactly-once delivery, which is a subtle but important distinction when evaluating a vendor's or a colleague's claim. At-least-once with a non-idempotent sink (a plain INSERT rather than an UPSERT) is a very common source of silent double-counting in feature pipelines, since the pipeline appears correct in testing (where failures are rare) and only manifests the bug under real-world transient failures in production.