KAFKA-20705: IQv2 TimestampedWindowKeyWithHeadersQuery for headers-aware window stores (KIP-1356)

streams/src/main/java/org/apache/kafka/streams/query/TimestampedWindowKeyWithHeadersQuery.java

@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.streams.query;
+
+import org.apache.kafka.common.annotation.InterfaceAudience;
+import org.apache.kafka.common.annotation.InterfaceStability.Evolving;
+import org.apache.kafka.streams.kstream.Windowed;
+import org.apache.kafka.streams.processor.api.ReadOnlyRecord;
+import org.apache.kafka.streams.state.ReadOnlyRecordIterator;
+import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders;
+
+import java.time.Instant;
+import java.util.Objects;
+import java.util.Optional;
+
+/**
+ * Interactive query for retrieving records of a single key across a window-start range, including
+ * their record headers, from a {@link TimestampedWindowStoreWithHeaders}.
+ *
+ * <p>This is the headers-aware parallel of {@link WindowKeyQuery}: it returns a
+ * {@link ReadOnlyRecordIterator} of {@link ReadOnlyRecord} elements, each carrying the windowed key,
+ * value, timestamp, and headers, whereas {@link WindowKeyQuery} returns a
+ * {@link org.apache.kafka.streams.state.WindowStoreIterator} of values keyed by window-start
+ * timestamp (value only, no headers). Each element's key is a {@link Windowed} whose window
+ * describes the record's window; the record's stored event-time is exposed via
+ * {@link ReadOnlyRecord#timestamp()}.
+ *
+ * <p>As with {@link WindowKeyQuery}, a closed window-start range must be supplied (both
+ * {@code timeFrom} and {@code timeTo}); open-ended ranges are not supported.
+ *
+ * <p>Headers are persisted and returned only when the store is backed by a native headers store,
+ * i.e. built with a KIP-1271 {@code WithHeaders} byte-store supplier (e.g.
+ * {@code Stores.persistentTimestampedWindowStoreWithHeaders}). A {@code WithHeaders} store built
+ * over a legacy (non-headers) supplier cannot persist headers, so the store-served reads come back
+ * with empty {@code headers()}.
+ *
+ * <p>Against a plain store not built with the {@code WithHeaders} builder at all, this query type is
+ * unsupported and fails with {@link FailureReason#UNKNOWN_QUERY_TYPE}.
+ *
+ * <p>Each element is a {@link ReadOnlyRecord}, whose timestamp is non-negative. If the backing store
+ * does not persist timestamps -- for example a {@code WithHeaders} store built over a plain
+ * window-store supplier that surfaces entries with {@code NO_TIMESTAMP} ({@code -1}) -- that entry
+ * cannot be represented, so advancing the returned {@link ReadOnlyRecordIterator} throws
+ * {@link org.apache.kafka.streams.errors.StreamsException} at that entry. Back the store with
+ * {@code Stores.persistentTimestampedWindowStoreWithHeaders(...)} to persist timestamps and headers,
+ * or use {@link WindowKeyQuery} if headers are not needed.
+ *
+ * <p>Because iteration can throw mid-stream, always close the returned {@link ReadOnlyRecordIterator}
+ * (for example with try-with-resources), even when a call to {@code next()} throws; otherwise the
+ * underlying store iterator leaks and the {@code num-open-iterators} metric stays incremented.
+ *
+ * @param <K> Type of keys
+ * @param <V> Type of values
+ */
+@Evolving
+@InterfaceAudience.Public
+public final class TimestampedWindowKeyWithHeadersQuery<K, V> implements Query<ReadOnlyRecordIterator<Windowed<K>, V>> {
+
+    private final K key;
+    private final Optional<Instant> timeFrom;
+    private final Optional<Instant> timeTo;
+
+    private TimestampedWindowKeyWithHeadersQuery(final K key,
+                                                 final Optional<Instant> timeFrom,
+                                                 final Optional<Instant> timeTo) {
+        this.key = key;
+        this.timeFrom = timeFrom;
+        this.timeTo = timeTo;
+    }
+
+    /**
+     * Creates a query that will retrieve the records (value, timestamp, and headers) identified by
+     * {@code key} whose window start falls within the closed range {@code [timeFrom, timeTo]}.
+     * @param key      The key to retrieve
+     * @param timeFrom The inclusive lower bound of the window-start range
+     * @param timeTo   The inclusive upper bound of the window-start range
+     * @param <K>      The type of the key
+     * @param <V>      The type of the value that will be retrieved
+     */
+    public static <K, V> TimestampedWindowKeyWithHeadersQuery<K, V> withKeyAndWindowStartRange(final K key,
+                                                                                               final Instant timeFrom,
+                                                                                               final Instant timeTo) {
+        Objects.requireNonNull(key, "the key should not be null");
+        return new TimestampedWindowKeyWithHeadersQuery<>(key, Optional.of(timeFrom), Optional.of(timeTo));
+    }
+
+    /**
+     * The key that was specified for this query.
+     */
+    public K key() {
+        return key;
+    }
+
+    /**
+     * The inclusive lower bound of the window-start range, if specified.
+     */
+    public Optional<Instant> timeFrom() {
+        return timeFrom;
+    }
+
+    /**
+     * The inclusive upper bound of the window-start range, if specified.
+     */
+    public Optional<Instant> timeTo() {
+        return timeTo;
+    }
+
+    @Override
+    public String toString() {
+        return "TimestampedWindowKeyWithHeadersQuery{" +
+            "key=" + key +
+            ", timeFrom=" + timeFrom +
+            ", timeTo=" + timeTo +
+            '}';
+    }
+}

streams/src/main/java/org/apache/kafka/streams/query/WindowKeyQuery.java

@@ -33,8 +33,8 @@ public class WindowKeyQuery<K, V> implements Query<WindowStoreIterator<V>> {
     private final Optional<Instant> timeTo;
 
     private WindowKeyQuery(final K key,
-                           final Optional<Instant> timeTo,
-                           final Optional<Instant> timeFrom) {
+                           final Optional<Instant> timeFrom,
+                           final Optional<Instant> timeTo) {
         this.key = key;
         this.timeFrom = timeFrom;
         this.timeTo = timeTo;

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

@@ -24,19 +24,25 @@
 import org.apache.kafka.common.utils.Time;
 import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.errors.ProcessorStateException;
+import org.apache.kafka.streams.errors.StreamsException;
 import org.apache.kafka.streams.kstream.Windowed;
+import org.apache.kafka.streams.kstream.internals.TimeWindow;
 import org.apache.kafka.streams.processor.StateStore;
+import org.apache.kafka.streams.processor.api.ReadOnlyRecord;
+import org.apache.kafka.streams.processor.api.Record;
 import org.apache.kafka.streams.processor.internals.ProcessorRecordContext;
 import org.apache.kafka.streams.processor.internals.SerdeGetter;
 import org.apache.kafka.streams.query.FailureReason;
 import org.apache.kafka.streams.query.PositionBound;
 import org.apache.kafka.streams.query.Query;
 import org.apache.kafka.streams.query.QueryConfig;
 import org.apache.kafka.streams.query.QueryResult;
+import org.apache.kafka.streams.query.TimestampedWindowKeyWithHeadersQuery;
 import org.apache.kafka.streams.query.WindowKeyQuery;
 import org.apache.kafka.streams.query.WindowRangeQuery;
 import org.apache.kafka.streams.query.internals.InternalQueryResultUtil;
 import org.apache.kafka.streams.state.KeyValueIterator;
+import org.apache.kafka.streams.state.ReadOnlyRecordIterator;
 import org.apache.kafka.streams.state.ReadOnlyWindowStore;
 import org.apache.kafka.streams.state.TimestampedBytesStore;
 import org.apache.kafka.streams.state.TimestampedWindowStoreWithHeaders;
@@ -142,6 +148,8 @@ public <R> QueryResult<R> query(
 
         if (query instanceof WindowKeyQuery) {
             result = runWindowKeyQuery((WindowKeyQuery<K, ValueTimestampHeaders<V>>) query, positionBound, config);
+        } else if (query instanceof TimestampedWindowKeyWithHeadersQuery) {
+            result = runTimestampedWindowKeyWithHeadersQuery((TimestampedWindowKeyWithHeadersQuery<K, V>) query, positionBound, config);
         } else if (query instanceof WindowRangeQuery) {
             result = runWindowRangeQuery((WindowRangeQuery<K, ValueTimestampHeaders<V>>) query, positionBound, config);
         } else {
@@ -216,6 +224,50 @@ private <R> QueryResult<R> runWindowKeyQuery(
         return queryResult;
     }
 
+    /**
+     * Handles {@link TimestampedWindowKeyWithHeadersQuery} by forwarding a raw byte-level
+     * {@link WindowKeyQuery} to the wrapped store and surfacing each entry as a {@link ReadOnlyRecord}
+     * (via {@link Record}) whose key is a {@link Windowed} of the queried key and the entry's window,
+     * carrying the stored value, event-time timestamp, and headers.
+     */
+    @SuppressWarnings("unchecked")
+    private <R> QueryResult<R> runTimestampedWindowKeyWithHeadersQuery(
+        final TimestampedWindowKeyWithHeadersQuery<K, V> query,
+        final PositionBound positionBound,
+        final QueryConfig config
+    ) {
+        final QueryResult<R> queryResult;
+
+        // Mirrors WindowKeyQuery handling: the closed-range contract (both bounds present) is guaranteed
+        // by the sole withKeyAndWindowStartRange factory today, but the Optional accessors are @Evolving,
+        // so we reject an open-ended range with UNKNOWN_QUERY_TYPE rather than assuming .get() is safe.
+        if (query.timeFrom().isPresent() && query.timeTo().isPresent()) {
+            final WindowKeyQuery<Bytes, byte[]> rawKeyQuery =
+                WindowKeyQuery.withKeyAndWindowStartRange(
+                    serializeKey(query.key(), internalContext.headers()),
+                    query.timeFrom().get(),
+                    query.timeTo().get()
+                );
+            final QueryResult<WindowStoreIterator<byte[]>> rawResult = wrapped().query(rawKeyQuery, positionBound, config);
+            if (rawResult.isSuccess()) {
+                final ReadOnlyRecordIterator<Windowed<K>, V> resultIterator =
+                    new MeteredWindowStoreWithHeadersReadOnlyRecordIterator(rawResult.getResult(), query.key());
+                queryResult = (QueryResult<R>) InternalQueryResultUtil.copyAndSubstituteDeserializedResult(rawResult, resultIterator);
+            } else {
+                queryResult = (QueryResult<R>) rawResult;
+            }
+        } else {
+            queryResult = QueryResult.forFailure(
+                FailureReason.UNKNOWN_QUERY_TYPE,
+                "This store (" + getClass() + ") doesn't know how to execute"
+                    + " the given query (" + query + ") because it only supports closed-range"
+                    + " queries."
+                    + " Contact the store maintainer if you need support for a new query type."
+            );
+        }
+        return queryResult;
+    }
+
     private <ValueType> MeteredWindowStoreIterator<ValueType> meteredIterator(
         final QueryResult<WindowStoreIterator<byte[]>> rawResult,
         final Function<byte[], ValueType> valueDeserializer
@@ -519,6 +571,92 @@ public Windowed<K> peekNextKey() {
         }
     }
 
+    /**
+     * Iterator backing {@link TimestampedWindowKeyWithHeadersQuery}: yields each entry as a
+     * {@link ReadOnlyRecord} (implemented by {@link Record}) whose key is a {@link Windowed} of the
+     * queried key and the entry's window, carrying the value, stored event-time timestamp, and the
+     * stored headers, with the headers frozen so a caller cannot mutate the read-only result.
+     *
+     * <p>A {@link ReadOnlyRecord} timestamp is contractually non-negative, so an entry with a
+     * negative stored timestamp cannot be represented -- for example a {@code WithHeaders} store
+     * built over a plain, non-timestamped window supplier surfaces entries with
+     * {@code NO_TIMESTAMP} (-1). This mirrors the rule the point/range headers queries apply. But
+     * because a lazily-evaluated iterator has already returned a successful {@link QueryResult}
+     * before any entry is read, such an entry cannot be surfaced as a query-level failure; it is
+     * instead reported by throwing a {@link StreamsException} while advancing the iterator.
+     *
+     * <p>Because {@link #next()} can throw mid-iteration, the caller must always close this iterator
+     * (via try-with-resources or a {@code finally} block) even when iteration throws: a caller that
+     * catches the exception and abandons the iterator without closing it leaks the underlying store
+     * iterator and leaves the {@code num-open-iterators} metric permanently incremented.
+     */
+    private class MeteredWindowStoreWithHeadersReadOnlyRecordIterator
+        implements ReadOnlyRecordIterator<Windowed<K>, V>, MeteredIterator {
+
+        private final WindowStoreIterator<byte[]> iter;
+        private final K key;
+        private final long startNs;
+        private final long startTimestampMs;
+
+        private MeteredWindowStoreWithHeadersReadOnlyRecordIterator(
+            final WindowStoreIterator<byte[]> iter,
+            final K key
+        ) {
+            this.iter = iter;
+            this.key = key;
+            this.startNs = time.nanoseconds();
+            this.startTimestampMs = time.milliseconds();
+            numOpenIterators.increment();
+            openIterators.add(this);
+        }
+
+        @Override
+        public long startTimestamp() {
+            return startTimestampMs;
+        }
+
+        @Override
+        public boolean hasNext() {
+            return iter.hasNext();
+        }
+
+        @Override
+        public ReadOnlyRecord<Windowed<K>, V> next() {
+            final KeyValue<Long, byte[]> next = iter.next();
+            final long windowStart = next.key;
+            final ValueTimestampHeaders<V> valueTimestampHeaders = deserializeValue(next.value);
+            final Headers headers = valueTimestampHeaders.headers();
+            if (valueTimestampHeaders.timestamp() < 0) {
+                throw new StreamsException(
+                    "Cannot represent the stored record for key [" + key + "] and window start ["
+                        + windowStart + "] as a ReadOnlyRecord: its timestamp ("
+                        + valueTimestampHeaders.timestamp() + ") is negative.");
+            }
+            final Windowed<K> windowedKey =
+                new Windowed<>(key, new TimeWindow(windowStart, windowStart + windowSizeMs));
+            final Record<Windowed<K>, V> record = new Record<>(
+                windowedKey,
+                valueTimestampHeaders.value(),
+                valueTimestampHeaders.timestamp(),
+                headers);
+            ((RecordHeaders) record.headers()).setReadOnly();
+            return record;
+        }
+
+        @Override
+        public void close() {
+            try {
+                iter.close();
+            } finally {
+                final long duration = time.nanoseconds() - startNs;
+                fetchSensor.record(duration);
+                iteratorDurationSensor.record(duration);
+                numOpenIterators.decrement();
+                openIterators.remove(this);
+            }
+        }
+    }
+
     private <ValueType> MeteredWindowedKeyValueIterator<K, ValueType> meteredWindowedIterator(
         final QueryResult<KeyValueIterator<Windowed<Bytes>, byte[]>> rawResult,
         final Function<ValueTimestampHeaders<V>, ValueType> valueConverter

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

@@ -67,7 +67,7 @@ public class MeteredWindowStore<K, V>
     extends WrappedStateStore<WindowStore<Bytes, byte[]>, Windowed<K>, V>
     implements WindowStore<K, V>, MeteredStateStore {
 
-    private final long windowSizeMs;
+    protected final long windowSizeMs;
     private final String metricsScope;
     protected final Time time;
     private final Serde<K> keySerde;

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

@@ -16,12 +16,12 @@
  */
 package org.apache.kafka.streams.state.internals;
 
-import org.apache.kafka.streams.processor.StateStore;
 import org.apache.kafka.streams.query.Position;
 import org.apache.kafka.streams.query.PositionBound;
 import org.apache.kafka.streams.query.Query;
 import org.apache.kafka.streams.query.QueryConfig;
 import org.apache.kafka.streams.query.QueryResult;
+import org.apache.kafka.streams.query.WindowKeyQuery;
 import org.apache.kafka.streams.state.HeadersBytesStore;
 import org.apache.kafka.streams.state.TimestampedBytesStore;
 
@@ -32,9 +32,6 @@
  * {@link TimestampedBytesStore} (for timestamp support) and {@link HeadersBytesStore}
  * (for header support) marker interfaces.
  * <p>
- * This store returns {@link QueryResult#forUnknownQueryType(Query, StateStore)} for all queries,
- * as IQv2 query handling is done at the metered layer.
- * <p>
  * The storage format for values is: [headersSize(varint)][headersBytes][timestamp(8)][value]
  *
  * @see RocksDBTimeOrderedWindowStore
@@ -53,6 +50,15 @@ class RocksDBTimeOrderedWindowStoreWithHeaders extends RocksDBTimeOrderedWindowS
     public <R> QueryResult<R> query(final Query<R> query,
                                     final PositionBound positionBound,
                                     final QueryConfig config) {
+        // KIP-1356 (window key query only): the headers-aware TimestampedWindowKeyWithHeadersQuery
+        // forwards a raw WindowKeyQuery to this native store, so enable WindowKeyQuery via the inherited
+        // RocksDBTimeOrderedWindowStore handling (StoreQueryUtils). Every other query type (e.g.
+        // WindowRangeQuery) stays UNKNOWN_QUERY_TYPE here, unchanged from before; those are enabled by
+        // their own KIP-1356 follow-ups.
+        if (query instanceof WindowKeyQuery) {
+            return super.query(query, positionBound, config);
+        }
+
         final long start = config.isCollectExecutionInfo() ? System.nanoTime() : -1L;
         final QueryResult<R> result;
         final Position position = getPosition();

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

@@ -21,6 +21,7 @@
 import org.apache.kafka.streams.query.Query;
 import org.apache.kafka.streams.query.QueryConfig;
 import org.apache.kafka.streams.query.QueryResult;
+import org.apache.kafka.streams.query.WindowKeyQuery;
 import org.apache.kafka.streams.state.HeadersBytesStore;
 import org.apache.kafka.streams.state.TimestampedBytesStore;
 
@@ -57,6 +58,15 @@ class RocksDBTimestampedWindowStoreWithHeaders extends RocksDBWindowStore implem
     public <R> QueryResult<R> query(final Query<R> query,
                                     final PositionBound positionBound,
                                     final QueryConfig config) {
+        // KIP-1356 (window key query only): the headers-aware TimestampedWindowKeyWithHeadersQuery
+        // forwards a raw WindowKeyQuery to this native store, so enable WindowKeyQuery via the inherited
+        // RocksDBWindowStore handling (StoreQueryUtils). Every other query type (e.g. WindowRangeQuery)
+        // stays UNKNOWN_QUERY_TYPE here, unchanged from before; those are enabled by their own KIP-1356
+        // follow-ups.
+        if (query instanceof WindowKeyQuery) {
+            return super.query(query, positionBound, config);
+        }
+
         final long start = config.isCollectExecutionInfo() ? System.nanoTime() : -1L;
         final QueryResult<R> result;
         final Position position = getPosition();