Mongo window range queries

kafka-client/src/main/java/dev/responsive/kafka/api/config/ResponsiveConfig.java

@@ -180,12 +180,13 @@ public class ResponsiveConfig extends AbstractConfig {
   // ------------------ WindowStore configurations ----------------------
 
   public static final String MONGO_WINDOWED_KEY_TIMESTAMP_FIRST_CONFIG = "responsive.mongo.windowed.key.timestamp.first";
-  private static final boolean MONGO_WINDOWED_KEY_TIMESTAMP_FIRST_DEFAULT = false;
+  private static final boolean MONGO_WINDOWED_KEY_TIMESTAMP_FIRST_DEFAULT = true;
   private static final String MONGO_WINDOWED_KEY_TIMESTAMP_FIRST_DOC = "Whether to put the window start timestamp "
       + "first in the composite windowed key format for MongoDB. This can be toggled true/false to get better "
       + "performance depending on the density of unique keys per window, and should be experimented "
-      + "with for best results. However it is important to note that this cannot be changed for "
-      + "an active application. Messing with this can corrupt existing state!";
+      + "with for best results. Must be true for any application that uses range queries on window stores. "
+      + "It is also important to note that this cannot be changed for an active application -- flipping "
+      + "this can corrupt existing state.";
 
   public static final String WINDOW_BLOOM_FILTER_COUNT_CONFIG = "responsive.window.bloom.filter.count";
   private static final int WINDOW_BLOOM_FILTER_COUNT_DEFAULT = 0;

kafka-client/src/main/java/dev/responsive/kafka/internal/db/CassandraWindowedTable.java

@@ -935,6 +935,22 @@ public KeyValueIterator<WindowedKey, byte[]> backFetchAll(
     );
   }
 
+  @Override
+  public KeyValueIterator<WindowedKey, byte[]> all(
+      final int kafkaPartition,
+      final long streamTime
+  ) {
+    throw new UnsupportedOperationException("all is not yet supported for Cassandra backends");
+  }
+
+  @Override
+  public KeyValueIterator<WindowedKey, byte[]> backAll(
+      final int kafkaPartition,
+      final long streamTime
+  ) {
+    throw new UnsupportedOperationException("backAll is not yet supported for Cassandra backends");
+  }
+
   private static KeyValue<WindowedKey, byte[]> windowRows(final Row row) {
     final long startTs = row.getInstant(WINDOW_START.column()).toEpochMilli();
     final Bytes key = Bytes.wrap(row.getByteBuffer(DATA_KEY.column()).array());

kafka-client/src/main/java/dev/responsive/kafka/internal/db/RemoteWindowedTable.java

@@ -18,6 +18,7 @@
 
 import dev.responsive.kafka.internal.utils.WindowedKey;
 import org.apache.kafka.common.utils.Bytes;
+import org.apache.kafka.streams.kstream.Windowed;
 import org.apache.kafka.streams.state.KeyValueIterator;
 
 public interface RemoteWindowedTable<S> extends RemoteTable<WindowedKey, S> {
@@ -157,4 +158,22 @@ KeyValueIterator<WindowedKey, byte[]> backFetchAll(
       long timeFrom,
       long timeTo
   );
+
+  /**
+   * Retrieves the windows of the given {@code kafkaPartition} across all keys and timestamps.
+   *
+   * @param kafkaPartition the kafka partition
+   *
+   * @return a forwards iterator over all the windows and values previously set
+   */
+  KeyValueIterator<WindowedKey, byte[]> all(int kafkaPartition, long streamTime);
+
+  /**
+   * Retrieves the windows of the given {@code kafkaPartition} across all keys and timestamps.
+   *
+   * @param kafkaPartition the kafka partition
+   *
+   * @return a backwards iterator over all the windows and values previously set
+   */
+  KeyValueIterator<WindowedKey, byte[]> backAll(int kafkaPartition, long streamTime);
 }

kafka-client/src/main/java/dev/responsive/kafka/internal/db/mongo/MongoWindowedTable.java

@@ -21,6 +21,7 @@
 import static com.mongodb.MongoClientSettings.getDefaultCodecRegistry;
 import static dev.responsive.kafka.internal.db.partitioning.SegmentPartitioner.UNINITIALIZED_STREAM_TIME;
 import static dev.responsive.kafka.internal.stores.ResponsiveStoreRegistration.NO_COMMITTED_OFFSET;
+import static dev.responsive.kafka.internal.stores.SegmentedOperations.MIN_KEY;
 import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
 import static org.bson.codecs.configuration.CodecRegistries.fromRegistries;
 
@@ -52,6 +53,7 @@
 import org.apache.kafka.common.utils.Bytes;
 import org.apache.kafka.streams.KeyValue;
 import org.apache.kafka.streams.errors.TaskMigratedException;
+import org.apache.kafka.streams.kstream.Windowed;
 import org.apache.kafka.streams.state.KeyValueIterator;
 import org.bson.codecs.configuration.CodecProvider;
 import org.bson.codecs.configuration.CodecRegistry;
@@ -457,13 +459,37 @@ public KeyValueIterator<WindowedKey, byte[]> backFetch(
   @Override
   public KeyValueIterator<WindowedKey, byte[]> fetchRange(
       final int kafkaPartition,
-      final Bytes fromKey,
-      final Bytes toKey,
+      final Bytes keyFrom,
+      final Bytes keyTo,
       final long timeFrom,
       final long timeTo
   ) {
-    throw new UnsupportedOperationException("fetchRange not yet supported for Mongo backends");
+    final List<KeyValueIterator<WindowedKey, byte[]>> segmentIterators = new LinkedList<>();
+    final var partitionSegments = kafkaPartitionToSegments.get(kafkaPartition);
+
+    for (final var segment : partitioner.range(kafkaPartition, timeFrom, timeTo)) {
+      final var segmentWindows = partitionSegments.segmentWindows.get(segment);
+
+      // Since we use a flat keyspace by concatenating the timestamp with the data key and have
+      // variable length data keys, it's impossible to request only valid data that's within
+      // the given bounds. Instead issue a broad request from the valid bounds and then post-filter
+      final var lowerBound = compositeKey(keyFrom, timeFrom);
+      final var upperBound = compositeKey(keyTo, timeTo);
 
+      final FindIterable<WindowDoc> fetchResults = segmentWindows.find(
+          Filters.and(
+              Filters.gte(WindowDoc.ID, lowerBound),
+              Filters.lte(WindowDoc.ID, upperBound))
+      );
+
+
+      segmentIterators.add(
+          Iterators.filterKv(
+              Iterators.kv(fetchResults.iterator(), MongoWindowedTable::windowFromDoc),
+              kv -> filterFetchRange(kv, timeFrom, timeTo, keyFrom, keyTo, timestampFirstOrder))
+      );
+    }
+    return Iterators.wrapped(segmentIterators);
   }
 
   @Override
@@ -483,7 +509,35 @@ public KeyValueIterator<WindowedKey, byte[]> fetchAll(
       final long timeFrom,
       final long timeTo
   ) {
-    throw new UnsupportedOperationException("fetchAll not yet supported for Mongo backends");
+    if (!timestampFirstOrder) {
+      throw new UnsupportedOperationException("Range queries such as fetchAll require stores to be "
+                                                  + "configured with timestamp-first order");
+    }
+
+    final List<KeyValueIterator<WindowedKey, byte[]>> segmentIterators = new LinkedList<>();
+    final var partitionSegments = kafkaPartitionToSegments.get(kafkaPartition);
+
+    for (final var segment : partitioner.range(kafkaPartition, timeFrom, timeTo)) {
+      final var segmentWindows = partitionSegments.segmentWindows.get(segment);
+
+      // To avoid scanning the entire segment, we use the bytewise "minimum key" to start the scan
+      // at the lower time bound. Since there's no corresponding "maximum key" given the variable
+      // length keys, we have to set the upper bound at timeTo + 1, while using strict comparison
+      // (ie #lt rather than #lte) to exclude said upper bound
+      final var lowerBound = compositeKey(MIN_KEY, timeFrom);
+      final var upperBound = compositeKey(MIN_KEY, timeTo + 1);
+
+      final FindIterable<WindowDoc> fetchResults = segmentWindows.find(
+          Filters.and(
+              Filters.gte(WindowDoc.ID, lowerBound),
+              Filters.lt(WindowDoc.ID, upperBound))
+      );
+
+      segmentIterators.add(
+          Iterators.kv(fetchResults.iterator(), MongoWindowedTable::windowFromDoc)
+      );
+    }
+    return Iterators.wrapped(segmentIterators);
   }
 
   @Override
@@ -495,6 +549,34 @@ public KeyValueIterator<WindowedKey, byte[]> backFetchAll(
     throw new UnsupportedOperationException("backFetchAll not yet supported for MongoDB backends");
   }
 
+  @Override
+  public KeyValueIterator<WindowedKey, byte[]> all(
+      final int kafkaPartition,
+      final long streamTime
+  ) {
+    final List<KeyValueIterator<WindowedKey, byte[]>> segmentIterators = new LinkedList<>();
+    final var partitionSegments = kafkaPartitionToSegments.get(kafkaPartition);
+
+    for (final var segment : partitioner.activeSegments(kafkaPartition, streamTime)) {
+      final var segmentWindows = partitionSegments.segmentWindows.get(segment);
+
+      final FindIterable<WindowDoc> fetchResults = segmentWindows.find();
+
+      segmentIterators.add(
+          Iterators.kv(fetchResults.iterator(), MongoWindowedTable::windowFromDoc)
+      );
+    }
+    return Iterators.wrapped(segmentIterators);
+  }
+
+  @Override
+  public KeyValueIterator<WindowedKey, byte[]> backAll(
+      final int kafkaPartition,
+      final long streamTime
+  ) {
+    throw new UnsupportedOperationException("backAll not yet supported for MongoDB backends");
+  }
+
   public BasicDBObject compositeKey(final WindowedKey windowedKey) {
     return compositeKey(windowedKey.key, windowedKey.windowStartMs);
   }
@@ -511,5 +593,22 @@ private static KeyValue<WindowedKey, byte[]> windowFromDoc(final WindowDoc windo
     return new KeyValue<>(WindowDoc.windowedKey(windowDoc.id), windowDoc.value);
   }
 
+  private static boolean filterFetchRange(
+      final WindowedKey windowedKey,
+      final long timeFrom,
+      final long timeTo,
+      final Bytes keyFrom,
+      final Bytes keyTo,
+      final boolean timestampFirstOrder
+  ) {
+    // If we use timestamp-first order, then the upper/lower bounds guarantee the timestamps are
+    // valid, so therefore we only need to filter out the invalid keys (and vice versa)
+    if (timestampFirstOrder) {
+      return windowedKey.key.compareTo(keyFrom) > 0 && windowedKey.key.compareTo(keyTo) < 0;
+    } else {
+      return windowedKey.windowStartMs > timeFrom && windowedKey.windowStartMs < timeTo;
+    }
+  }
+
 }
 

kafka-client/src/main/java/dev/responsive/kafka/internal/stores/CommitBuffer.java

@@ -342,6 +342,20 @@ public KeyValueIterator<K, Result<K>> range(final K from, final K to) {
     );
   }
 
+  public KeyValueIterator<K, Result<K>> range(
+      final K from,
+      final K to,
+      final boolean fromInclusive,
+      final boolean toInclusive
+  ) {
+    return Iterators.kv(
+        Iterators.filter(
+            buffer.getReader().subMap(from, fromInclusive, to, toInclusive).entrySet().iterator(),
+            e -> keySpec.retain(e.getKey())),
+        result -> new KeyValue<>(result.getKey(), result.getValue())
+    );
+  }
+
   public KeyValueIterator<K, Result<K>> range(
       final K from,
       final K to,

kafka-client/src/main/java/dev/responsive/kafka/internal/stores/ResponsiveWindowStore.java

@@ -270,7 +270,7 @@ public KeyValueIterator<Windowed<Bytes>, byte[]> fetchAll(
 
   @Override
   public KeyValueIterator<Windowed<Bytes>, byte[]> all() {
-    return windowOperations.all();
+    return windowOperations.all(observedStreamTime);
   }
 
   @Override

kafka-client/src/main/java/dev/responsive/kafka/internal/stores/SegmentedOperations.java

@@ -55,6 +55,9 @@
 
 public class SegmentedOperations implements WindowOperations {
 
+  // The "minimum" possible key when comparing bytewise, used to define range query bounds
+  public static final Bytes MIN_KEY = Bytes.wrap(new byte[0]);
+
   @SuppressWarnings("rawtypes")
   private final InternalProcessorContext context;
   private final ResponsiveWindowParams params;
@@ -267,12 +270,26 @@ public KeyValueIterator<Windowed<Bytes>, byte[]> fetchAll(
       final long timeFrom,
       final long timeTo
   ) {
-    throw new UnsupportedOperationException("Not yet implemented");
+    final WindowedKey lowerBound = new WindowedKey(MIN_KEY, timeFrom);
+    final WindowedKey upperBound = new WindowedKey(MIN_KEY, timeTo + 1);
+
+    // set toInclusive bound to false due to the +1 in the upper bound
+    return Iterators.windowedKey(
+        new LocalRemoteKvIterator<>(
+            buffer.range(lowerBound, upperBound, true, false),
+            table.fetchAll(changelog.partition(), timeFrom, timeTo)),
+        params.windowSize()
+    );
   }
 
   @Override
-  public KeyValueIterator<Windowed<Bytes>, byte[]> all() {
-    throw new UnsupportedOperationException("Not yet implemented");
+  public KeyValueIterator<Windowed<Bytes>, byte[]> all(final long streamTime) {
+    return Iterators.windowedKey(
+        new LocalRemoteKvIterator<>(
+            buffer.all(),
+            table.all(changelog.partition(), streamTime)),
+        params.windowSize()
+    );
   }
 
   @Override

kafka-client/src/main/java/dev/responsive/kafka/internal/stores/WindowOperations.java

@@ -53,7 +53,7 @@ KeyValueIterator<Windowed<Bytes>, byte[]> fetchAll(
       final long timeTo
   );
 
-  KeyValueIterator<Windowed<Bytes>, byte[]> all();
+  KeyValueIterator<Windowed<Bytes>, byte[]> all(final long streamTime);
 
   WindowStoreIterator<byte[]> backwardFetch(
       final Bytes key,