KAFKA-20318: Add time ordered session store support to the DSL by bbejeck · Pull Request #21794 · apache/kafka

bbejeck · 2026-03-17T19:20:34Z

This PR fixes a bug where
RocksDbTimeOrderedSessionBytesStoreSupplier.get with
withHeaders=true created segments using KeyValueSegment (default-CF
only) instead of SessionSegmentWithHeaders (dual-CF with lazy
migration), which would cause the upgrade path from non-headers to
headers format to fail. The fix introduces
RocksDBTimeOrderedSessionSegmentedBytesStoreWithHeaders, widens
RocksDBTimeOrderedSessionStore to accept either segment type, and
updates the supplier to create the correct bytes store. Upgrade tests
validate lazy migration at both the store layer and DSL supplier path,
plus a TopologyTestDriver end-to-end test exercises ON_WINDOW_CLOSE
with both PLAIN and HEADERS formats.

Reviewers: Matthias J. Sax matthias@confluent.io, TengYao Chi
frankvicky@apache.org

bbejeck · 2026-03-17T19:26:58Z

.../org/apache/kafka/streams/state/internals/AbstractRocksDBTimeOrderedSegmentedBytesStore.java

+        throw new UnsupportedOperationException("This store does not support put with timestamp and seqnum");
+    }
+
+    public byte[] fetch(final Bytes key, final long timestamp, final int seqnum) {


Added fetchSession, fetchSessions, remove(Windowed), and put(Windowed) as UnsupportedOperationException defaults to match the existing pattern for put(Bytes,long,int,byte[]) and fetch(Bytes,long,int). This allows
RocksDBTimeOrderedSessionStore to widen its wrapped type to the abstract class while still calling session-specific methods. Both RocksDBTimeOrderedSessionSegmentedBytesStore and the new WithHeaders variant override these.

bbejeck · 2026-03-17T19:28:05Z

...s/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedSessionStore.java


 public class RocksDBTimeOrderedSessionStore
-    extends WrappedStateStore<RocksDBTimeOrderedSessionSegmentedBytesStore, Object, Object>
+    extends WrappedStateStore<AbstractRocksDBTimeOrderedSegmentedBytesStore<? extends Segment>, Object, Object>


Widened from WrappedStateStore<RocksDBTimeOrderedSessionSegmentedBytesStore, ...> to WrappedStateStore<AbstractRocksDBTimeOrderedSegmentedBytesStore<? extends Segment>, ...> so that both the plain and WithHeaders segmented
bytes stores can be wrapped. All wrapped() calls resolve against the abstract type which now has the session-specific methods from the previous change.

bbejeck · 2026-03-17T19:30:38Z

...he/kafka/streams/state/internals/RocksDBTimeOrderedWindowSegmentedBytesStoreWithHeaders.java

+ * @see WindowSegmentsWithHeaders
+ * @see WindowSegmentWithHeaders
+ */
+class RocksDBTimeOrderedWindowSegmentedBytesStoreWithHeaders


Mirrors RocksDBTimeOrderedSessionSegmentedBytesStoreWithHeaders after RocksDBTimeOrderedWindowSegmentedBytesStoreWithHeaders (PR #21780) but for session key schemas. Uses SessionSegmentsWithHeaders which creates
SessionSegmentWithHeaders segments — these extend RocksDBMigratingSessionStoreWithHeaders and handle dual-CF lazy migration from DEFAULT CF (old format) to headers CF (new format) via HeadersBytesStore::convertToHeaderFormat.
No new segment classes needed — reuses existing infrastructure.

bbejeck · 2026-03-17T19:35:49Z

...va/org/apache/kafka/streams/state/internals/RocksDbTimeOrderedSessionBytesStoreSupplier.java


    @Override
    public SessionStore<Bytes, byte[]> get() {
+        if (withHeaders) {


Previously, get() always created a RocksDBTimeOrderedSessionSegmentedBytesStore (which uses KeyValueSegment / DEFAULT CF only), even when withHeaders=true. On upgrade from non-headers to headers, old
data in DEFAULT CF would not be readable because KeyValueSegment has no migration logic. Now the withHeaders path creates RocksDBTimeOrderedSessionSegmentedBytesStoreWithHeaders which uses SessionSegmentWithHeaders segments
with dual-CF migration support.

bbejeck · 2026-03-17T19:37:12Z

...ms/src/test/java/org/apache/kafka/streams/state/internals/AbstractSessionBytesStoreTest.java

-                                        name(), metricsScope(), retentionPeriod(), segmentIntervalMs(), true));
-                            }
-                        },
+                        new RocksDbTimeOrderedSessionBytesStoreSupplier(ROCK_DB_STORE_NAME, retentionPeriod, true, true),


Removed the anonymous RocksDbTimeOrderedSessionBytesStoreSupplier override hack that was manually constructing a plain RocksDBTimeOrderedSessionSegmentedBytesStore for the WithHeaders test cases. Now uses the supplier's real
withHeaders=true code path directly.

Same for below

bbejeck · 2026-03-17T19:53:53Z

...c/test/java/org/apache/kafka/streams/state/internals/TimeOrderedSessionStoreUpgradeTest.java

+ * This test validates lazy migration from DEFAULT CF to headers CF and ensures
+ * the store can read old data after upgrade.
+ */
+public class TimeOrderedSessionStoreUpgradeTest {


Three layers of upgrade testing: (1) Store-level tests validate lazy migration directly using RocksDbTimeOrderedSessionBytesStoreSupplier — write with withHeaders=false, close, reopen with withHeaders=true, verify old data
readable with empty headers. (2) DSL supplier path test uses BuiltInDslStoreSuppliers.RocksDBDslStoreSuppliers.sessionStore() — the same code path the DSL materializer calls with ON_WINDOW_CLOSE. (3) DSL end-to-end test builds
a real session aggregation topology via TopologyTestDriver with ON_WINDOW_CLOSE + DSL_STORE_FORMAT_HEADERS, verifying the full chain works.

mjsax · 2026-03-19T02:57:55Z

.../org/apache/kafka/streams/state/internals/AbstractRocksDBTimeOrderedSegmentedBytesStore.java

 */
 public abstract class AbstractRocksDBTimeOrderedSegmentedBytesStore<S extends Segment> extends AbstractDualSchemaRocksDBSegmentedBytesStore<S> {
-    private static final Logger LOG = LoggerFactory.getLogger(AbstractDualSchemaRocksDBSegmentedBytesStore.class);
+    private static final Logger LOG = LoggerFactory.getLogger(AbstractRocksDBTimeOrderedSegmentedBytesStore.class);


mjsax · 2026-03-19T02:59:19Z

.../org/apache/kafka/streams/state/internals/AbstractRocksDBTimeOrderedSegmentedBytesStore.java

+    private static final Logger LOG = LoggerFactory.getLogger(AbstractRocksDBTimeOrderedSegmentedBytesStore.class);

-    abstract class IndexToBaseStoreIterator implements KeyValueIterator<Bytes, byte[]> {
+    private long minTimestamp;


I unified more code (cf other comments) -- this is moved from RocksDBTimeOrderedKeyValueBytesStore to here

mjsax · 2026-03-19T02:59:57Z

.../org/apache/kafka/streams/state/internals/AbstractRocksDBTimeOrderedSegmentedBytesStore.java

-    abstract class IndexToBaseStoreIterator implements KeyValueIterator<Bytes, byte[]> {
+    private long minTimestamp;
+
+    public abstract class IndexToBaseStoreIterator implements KeyValueIterator<Bytes, byte[]> {


Another side cleanup -- we use many of these classes outside of the internals package, so we should just declare them public.

There is more changes like this on this PR -- if we think we don't want this change, we can also revert it.

mjsax · 2026-03-19T03:00:29Z

.../org/apache/kafka/streams/state/internals/AbstractRocksDBTimeOrderedSegmentedBytesStore.java

-     * <p>
-     * This can be reused by both window store implementations (with and without headers).
-     */
-    class WindowKeySchemaIndexToBaseStoreIterator extends IndexToBaseStoreIterator {


This class belongs to RocksDBTimeOrderedWindowSegmentedBytesStore -- moved it there

mjsax · 2026-03-19T03:02:01Z

.../org/apache/kafka/streams/state/internals/AbstractRocksDBTimeOrderedSegmentedBytesStore.java

+        minTimestamp = Long.MAX_VALUE;
+    }
+
+    Map<S, WriteBatch> getWriteBatches(


This is a refactoring to allow sharing more code. The different sub-classes of AbstractRocksDBTimeOrderedSegmentedBytesStore all share a 99% common impl of Map<S, WriteBatch> getWriteBatches(Collection<ConsumerRecord<byte[], byte[]>> records), so I add this shared impl which is "customized" by passing corresponding extractors.

mjsax · 2026-03-19T03:07:12Z

streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegments.java

-    LogicalKeyValueSegment getReservedSegment(final long segmentId) {
-        return reservedSegments.get(segmentId);
+    LogicalKeyValueSegment getReservedSegment() {
+        return reservedSegments.get(-1L);


Side cleanup -- we always pas -1L.

mjsax · 2026-03-19T03:07:32Z

...s/src/main/java/org/apache/kafka/streams/state/internals/RocksDBSessionStoreWithHeaders.java

 * <p>
 * This store extends {@link RocksDBSessionStore} and returns
- * {@link QueryResult#forUnknownQueryType(Query, Object)} for all queries,
+ * {@link QueryResult#forUnknownQueryType(Query, StateStore)} for all queries,


Side fix (also elsewhere)

mjsax · 2026-03-19T03:08:52Z

...e/kafka/streams/state/internals/RocksDBTimeOrderedSessionSegmentedBytesStoreWithHeaders.java

+ * @see SessionSegmentsWithHeaders
+ * @see SessionSegmentWithHeaders
+ */
+class RocksDBTimeOrderedSessionSegmentedBytesStoreWithHeaders


Added the new class we need is not trivial with the unified code on the super-classes :)

mjsax · 2026-03-19T03:10:09Z

...ms/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedWindowStore.java


-public class RocksDBTimeOrderedWindowStore
-    extends WrappedStateStore<AbstractRocksDBTimeOrderedSegmentedBytesStore<? extends Segment>, Object, Object>
+public class RocksDBTimeOrderedWindowStore<S extends Segment>


This the the follow up from #21780

mjsax · 2026-03-19T03:10:41Z

.../java/org/apache/kafka/streams/state/internals/TimestampedWindowStoreWithHeadersBuilder.java

        if (stateStore instanceof RocksDBTimeOrderedWindowStore) {
            return true;
        }
-        if (stateStore instanceof RocksDBTimeOrderedWindowStoreWithHeaders) {


RocksDBTimeOrderedWindowStoreWithHeaders extends RocksDBTimeOrderedWindowStore so this second check is unnecessary.

Great catch

frankvicky · 2026-03-19T14:58:58Z

...s/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimeOrderedSessionStore.java

-    RocksDBTimeOrderedSessionStore(final RocksDBTimeOrderedSessionSegmentedBytesStore store) {
+    RocksDBTimeOrderedSessionStore(final AbstractRocksDBTimeOrderedSegmentedBytesStore<? extends Segment> store) {
        super(store);
        Objects.requireNonNull(store, "store is null");


Just a random things come to my mind:
There are are bunch of ternary operator for null check, do we consider to replace it with Objects.requireNonNullElse to align the code style?

This PR fixes a bug where `RocksDbTimeOrderedSessionBytesStoreSupplier.get` with `withHeaders=true` created segments using `KeyValueSegment` (default-CF only) instead of `SessionSegmentWithHeaders` (dual-CF with lazy migration), which would cause the upgrade path from non-headers to headers format to fail. The fix introduces `RocksDBTimeOrderedSessionSegmentedBytesStoreWithHeaders`, widens `RocksDBTimeOrderedSessionStore` to accept either segment type, and updates the supplier to create the correct bytes store. Upgrade tests validate lazy migration at both the store layer and DSL supplier path, plus a `TopologyTestDriver` end-to-end test exercises ON_WINDOW_CLOSE with both PLAIN and HEADERS formats. Reviewers: Matthias J. Sax <matthias@confluent.io>, TengYao Chi <frankvicky@apache.org> Co-authored-by: Matthias J. Sax <matthias@confluent.io>

mjsax · 2026-03-19T18:35:14Z

Merged to trunk and cherry-picked to 4.3 branch.

github-actions bot added the streams label Mar 17, 2026

bbejeck commented Mar 17, 2026

View reviewed changes

PR enables time-ordered session store with headers

859ccdf

mjsax added the kip Requires or implements a KIP label Mar 19, 2026

mjsax force-pushed the KAFKA-20318_add_time_ordered_session_store_dsl branch from 66a93c8 to 859ccdf Compare March 19, 2026 00:58

mjsax added 2 commits March 18, 2026 18:19

improve generics (follow up from PR#21780)

a6935dc

Refactoring

02b8806

mjsax mentioned this pull request Mar 19, 2026

KAFKA-20317: Enable time-ordered header window store in DSL #21780

Merged

cleanup

3fab7a3

mjsax reviewed Mar 19, 2026

View reviewed changes

frankvicky reviewed Mar 19, 2026

View reviewed changes

frankvicky approved these changes Mar 19, 2026

View reviewed changes

mjsax merged commit 1ad463d into apache:trunk Mar 19, 2026
26 checks passed

Conversation

bbejeck commented Mar 17, 2026 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bbejeck Mar 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mjsax Mar 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mjsax Mar 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mjsax commented Mar 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bbejeck commented Mar 17, 2026 •

edited by github-actions bot

Loading

bbejeck Mar 17, 2026 •

edited

Loading

mjsax Mar 19, 2026 •

edited

Loading

mjsax Mar 19, 2026 •

edited

Loading