KAFKA-20249: Optimize raw value extraction in headers-aware deserializers by zheguang · Pull Request #21706 · apache/kafka

zheguang · 2026-03-11T08:29:02Z

This patch implements two optimizations, and their JMH benchmarks, and
opportunistic refactoring.

Skipping
Previously the raw value extraction in headers-aware deserializers
undergoes deserialization and/or copying of headers, while only skipping
is required. This happens for both empty and nonempty headers.
Empty headers copying
Empty headers have constant metadata footprint: the headers size is
varint-encoded 1 byte of 0, and headers themselves consume no bytes.
Based on this invariant, the ByteBuffer-based extraction can be replaced
with a direct System.arraycopy, which is a Java native method
optimized for specific platforms.

The optimized headers-aware extraction methods:

rawAggregation
rawTimestampedValue
rawValue / rawPlainValue

Benchmark:
This patch also includes JMH benchmarks to test the speedup. On my
local machine, Optimization 1 speedup is 2-6x speedup. Optimization 2
is 1.2-1.3x.

Below is the throughput comparison of a recorded JMH run (higher score
is better):

Benchmark
Mode  Cnt      Score      Error  Units
RawBytesExtractionBenchmark.testHeadersWithoutHeaders
thrpt   15  10158.764 ±   85.564  ops/s
RawBytesExtractionBenchmark.testHeadersWithoutHeadersOpt
thrpt   15  14824.176 ± 1244.455  ops/s

RawBytesExtractionBenchmark.testRawAggregationWithHeaders
thrpt   15   1473.459 ±    7.170  ops/s
RawBytesExtractionBenchmark.testRawAggregationWithHeadersOpt
thrpt   15  11618.187 ±  235.385  ops/s

RawBytesExtractionBenchmark.testRawAggregationWithoutHeaders
thrpt   15   8337.728 ±  199.919  ops/s
RawBytesExtractionBenchmark.testRawAggregationWithoutHeadersOpt
thrpt   15  14564.899 ±  186.405  ops/s

RawBytesExtractionBenchmark.testRawTimestampedValueWithoutHeaders
thrpt   15  10217.292 ±  108.552  ops/s
RawBytesExtractionBenchmark.testRawTimestampedValueWithoutHeadersOpt
thrpt   15  12121.074 ±  201.235  ops/s

RawBytesExtractionBenchmark.testRawValueWithoutHeaders
thrpt   15  11632.484 ±  138.505  ops/s
RawBytesExtractionBenchmark.testRawValueWithoutHeadersOpt
thrpt   15  14669.563 ±   43.458  ops/s

RawBytesExtractionBenchmark.testTimestampWithoutHeaders
thrpt   15  14858.778 ±   39.301  ops/s
RawBytesExtractionBenchmark.testTimestampWithoutHeadersOpt
thrpt   15  19832.718 ±  916.980  ops/s
JMH benchmarks done

Test:

AggregationWithHeadersDeserializer.rawAggregate
- empty headers:
  SessionToHeadersStoreAdapterTest.shouldStripHeadersFromRawAggregationValue
Utils.rawPlainValue
- empty headers: UtilsTest.shouldExtractRawValueWithEmptyHeaders
- empty headers, no timestamp:
  UtilsTest.testRawPlainValueWithEmptyHeadersAndInvalidTimestamp
Utils.rawTimestampedValue
- empty headers: UtilsTest.testRawTimestampedValueWithEmptyHeaders
- empty headers, no timestamp:
  UtilsTest.testRawTimestampedValueWithEmptyHeadersAndInvalidTimestamp

Refactor

point all calls to raw value (with timestamp and headers) extraction
to common one in Utils.

Benchmark Mode Cnt Score Error Units RawBytesExtraction.testRawAggregationWithoutHeaders thrpt 15 7850.891 ± 307.428 ops/s RawBytesExtraction.testRawAggregationWithoutHeadersFastPath thrpt 15 14957.556 ± 517.450 ops/s

Benchmark Mode Cnt Score Error Units RawBytesExtraction.testRawAggregationWithHeaders thrpt 15 1411.338 ± 110.527 ops/s RawBytesExtraction.testRawAggregationWithHeadersFastPath thrpt 15 6106.665 ± 218.032 ops/s RawBytesExtraction.testRawAggregationWithoutHeaders thrpt 15 7734.538 ± 525.487 ops/s RawBytesExtraction.testRawAggregationWithoutHeadersFastPath thrpt 15 14300.408 ± 212.519 ops/s

Benchmark Mode Cnt Score Error Units RawBytesExtraction.testRawAggregationWithHeaders thrpt 15 1481.854 ± 31.448 ops/s RawBytesExtraction.testRawAggregationWithHeadersOpt thrpt 15 11797.165 ± 103.432 ops/s RawBytesExtraction.testRawAggregationWithoutHeaders thrpt 15 8359.080 ± 47.918 ops/s RawBytesExtraction.testRawAggregationWithoutHeadersOpt thrpt 15 15298.827 ± 452.741 ops/s RawBytesExtraction.testRawValueWithoutHeaders thrpt 15 11329.997 ± 260.399 ops/s RawBytesExtraction.testRawValueWithoutHeadersOpt thrpt 15 15372.816 ± 184.651 ops/s

zheguang · 2026-03-12T23:44:02Z

Hi @aliehsaeedii - if you could please have a quick look and see if this approach is headed in the right direction? Thanks!

aliehsaeedii

Thanks @zheguang for the PR. Could we apply this optimization (does that make sense to apply it) in HeadersDeserializer, TimestampedToHeadersWindowStoreAdapter, and HeadersBytesStore classes?
Please add more utests so that all changes are tested.

aliehsaeedii · 2026-03-13T17:36:49Z

...c/main/java/org/apache/kafka/streams/state/internals/AggregationWithHeadersDeserializer.java

 * This is used by KIP-1271 to deserialize aggregations with headers from session state stores.
 */
-class AggregationWithHeadersDeserializer<AGG> implements WrappingNullableDeserializer<AggregationWithHeaders<AGG>, Void, AGG> {
+public class AggregationWithHeadersDeserializer<AGG> implements WrappingNullableDeserializer<AggregationWithHeaders<AGG>, Void, AGG> {


I assume you make them public to be used in jmh testing?!

That's right. It's a bit awkward to broaden this deserializer's scope, tbh. One option is to move those benchmarked methods out of this deserializer, into Utils, and instead only broaden the scope of Utils... Would this be preferred? Let me know.

I'm not sure if we need to keep the jmh benchmarks!

aliehsaeedii · 2026-03-13T17:48:00Z

...c/main/java/org/apache/kafka/streams/state/internals/AggregationWithHeadersDeserializer.java

+
        final ByteBuffer buffer = ByteBuffer.wrap(aggregationWithHeaders);
-        readHeaders(buffer);
+        // Skip the headers bytes without deserizization or copying


Suggested change

// Skip the headers bytes without deserizization or copying

// Skip the headers bytes without deserialization or copying

aliehsaeedii · 2026-03-13T17:56:46Z

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

        return result;
    }

+    private static boolean hasEmptyHeadersAndTimestamp(final byte[] rawValueTimestampHeaders) {


nit: Should it be hasEmptyHeaders only? We don't check empty ts!

Actualy... it does check a little something about the timestamp -- that the input is at least longer than the timestamp size. Line 64 is the relevant bit:

if (rawValueTimestampHeaders.length - 1 - StateSerdes.TIMESTAMP_SIZE < 0)

aliehsaeedii · 2026-03-13T17:58:42Z

streams/src/test/java/org/apache/kafka/streams/state/internals/UtilsTest.java

+        final ByteBuffer buf = ByteBuffer.wrap(data);
+        buf.put((byte) 0x00); // header size
+        buf.putLong(TIMESTAMP);
+        buf.put(VALUE); // non-header payload


Suggested change

buf.put(VALUE); // non-header payload

buf.put(VALUE); // plain value

aliehsaeedii · 2026-03-13T18:09:48Z

...c/main/java/org/apache/kafka/streams/state/internals/AggregationWithHeadersDeserializer.java

    }

-    private static Headers readHeaders(final ByteBuffer buffer) {
+    public static Headers readHeaders(final ByteBuffer buffer) {


can we optimize readHeaders as well?

Yea, fast path for empty headers gives 1.5x speedup on my local machine

Benchmark Mode Cnt Score Error Units RawBytesExtractionBenchmark.testHeadersWithoutHeaders thrpt 15 10198.854 ± 62.216 ops/s RawBytesExtractionBenchmark.testHeadersWithoutHeadersOpt thrpt 15 15852.852 ± 47.469 ops/s

I'll add this change to this PR.

You mean headers() (which calls readHeaders)? Yes I can see a speedup too:

RawBytesExtractionBenchmark.testHeadersWithoutHeaders thrpt 15 10158.764 ± 85.564 ops/s RawBytesExtractionBenchmark.testHeadersWithoutHeadersOpt thrpt 15 14824.176 ± 1244.455 ops/s

I will make this change in this PR.

aliehsaeedii · 2026-03-13T20:33:23Z

...c/main/java/org/apache/kafka/streams/state/internals/AggregationWithHeadersDeserializer.java

+    public static Headers readHeaders(final ByteBuffer buffer) {
        final int headersSize = ByteUtils.readVarint(buffer);
        final byte[] rawHeaders = readBytes(buffer, headersSize);
        return HeadersDeserializer.deserialize(rawHeaders);


I'm wondering if it makes sense to do the if (rawAggregationWithHeaders.length > 0 && rawAggregationWithHeaders[0] == 0x00) { in static Headers headers(final byte[] rawAggregationWithHeaders) as well!

Hm... not sure -- for empty headers there is a fast (enough?) path in HeadersDeserializer.deserialize:

// in HeadersDeserializer public static Headers deserialize(final byte[] data) { if (data == null || data.length == 0) { return new RecordHeaders(); }

I'll find out in the context of KAFKA-20303 to be sure though.

aliehsaeedii · 2026-03-13T20:41:20Z

...rc/main/java/org/apache/kafka/streams/state/internals/ValueTimestampHeadersDeserializer.java

-        final int headersSize = ByteUtils.readVarint(buffer);
-        buffer.position(buffer.position() + headersSize + Long.BYTES);
-        return readBytes(buffer, buffer.remaining());
-    }


Does that make sense to apply the same optimization in the methods of the class such as headers(), value(), and deserialize()?

Great pointer. value() indeed can just call rawPlainValue(). So I just added this change in this PR.

~~For the others, let me find out in the context of KAFKA-20303~~

I made the change in this PR anyways to find out ... yes there is a speedup to headers and timestsamp for empty headers:

RawBytesExtractionBenchmark.testHeadersWithoutHeaders thrpt 15 10158.764 ± 85.564 ops/s RawBytesExtractionBenchmark.testHeadersWithoutHeadersOpt thrpt 15 14824.176 ± 1244.455 ops/s RawBytesExtractionBenchmark.testTimestampWithoutHeaders thrpt 15 14858.778 ± 39.301 ops/s RawBytesExtractionBenchmark.testTimestampWithoutHeadersOpt thrpt 15 19832.718 ± 916.980 ops/s

aliehsaeedii · 2026-03-13T20:53:04Z

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

+        if (hasEmptyHeadersAndTimestamp(rawValueTimestampHeaders)) {
+            // Strip header size (varint 1 byte), empty headers (no bytes), and timestamp
+            final byte[] res = new byte[rawValueTimestampHeaders.length - 1 - StateSerdes.TIMESTAMP_SIZE]; 
+            System.arraycopy(rawValueTimestampHeaders, 1 + StateSerdes.TIMESTAMP_SIZE, res, 0, res.length);


Should validate rawValueTimestampHeaders.length >= 1 + StateSerdes.TIMESTAMP_SIZE

Yep, this is validated within the hasEmptyHeadersAndTimestamp call above, at l.64:

if (rawValueTimestampHeaders.length - 1 - StateSerdes.TIMESTAMP_SIZE < 0) { // throw serialization exception

zheguang added 4 commits March 11, 2026 14:02

Performance test of fast path raw value bytes extraction

e44a2c5

Benchmark Mode Cnt Score Error Units RawBytesExtraction.testRawAggregationWithoutHeaders thrpt 15 7850.891 ± 307.428 ops/s RawBytesExtraction.testRawAggregationWithoutHeadersFastPath thrpt 15 14957.556 ± 517.450 ops/s

Fix styles

754763e

github-actions bot added triage PRs from the community streams performance labels Mar 11, 2026

zheguang added 5 commits March 12, 2026 08:34

Merge remote-tracking branch 'origin/trunk' into zheguang-KAFKA-20249

1c65c8c

Fix merge error

83998cd

Add boundary check for empty headers

b761933

Refactor rawValue to Utils

4d9a42f

Optimize raw value with timestamp extraction

1e26a88

frankvicky added the ci-approved label Mar 13, 2026

aliehsaeedii reviewed Mar 13, 2026

View reviewed changes

github-actions bot removed the triage PRs from the community label Mar 14, 2026

zheguang added 4 commits March 14, 2026 22:12

Add unit tests for invalid data (no timestamp)

da95328

Point ValueTimestampHeadersDeserializer to Utils.rawPlainValue

af532a2

Optimize aggregate deserializer's header extraction

d4d4260

Optimize timestamp extraction

aa6277f

	// Skip the headers bytes without deserizization or copying
	// Skip the headers bytes without deserialization or copying

	buf.put(VALUE); // non-header payload
	buf.put(VALUE); // plain value

Conversation

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

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

aliehsaeedii left a comment • 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!

aliehsaeedii Mar 14, 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!

zheguang Mar 14, 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!

zheguang Mar 14, 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!

zheguang Mar 14, 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!

aliehsaeedii Mar 13, 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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zheguang commented Mar 11, 2026 •

edited by github-actions bot

Loading

zheguang commented Mar 12, 2026 •

edited

Loading

aliehsaeedii left a comment •

edited

Loading

aliehsaeedii Mar 14, 2026 •

edited

Loading

zheguang Mar 14, 2026 •

edited

Loading

zheguang Mar 14, 2026 •

edited

Loading

zheguang Mar 14, 2026 •

edited

Loading

aliehsaeedii Mar 13, 2026 •

edited

Loading