dev.responsive.kafka.internal.utils.Iterators Maven / Gradle / Ivy
The newest version!
/*
* Copyright 2024 Responsive Computing, Inc.
*
* This source code is licensed under the Responsive Business Source License Agreement v1.0
* available at:
*
* https://www.responsive.dev/legal/responsive-bsl-10
*
* This software requires a valid Commercial License Key for production use. Trial and commercial
* licenses can be obtained at https://www.responsive.dev
*/
package dev.responsive.kafka.internal.utils;
import static java.util.Collections.emptyIterator;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.function.Function;
import java.util.function.Predicate;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.streams.KeyValue;
import org.apache.kafka.streams.kstream.Windowed;
import org.apache.kafka.streams.kstream.internals.TimeWindow;
import org.apache.kafka.streams.state.KeyValueIterator;
import org.apache.kafka.streams.state.WindowStoreIterator;
/**
* Utility class for creating iterators.
*/
public final class Iterators {
private Iterators() {}
public static KeyValueIterator emptyKv() {
return Iterators.kv(emptyIterator(), e -> new KeyValue<>(null, null));
}
/**
* Returns an iterator that caches the last value returned by
* a delegate so that {@code peek()} can check the next key
* without exhausting the iterator.
*
* @param delegate the delegate
* @return a new iterator of the same type as delegate that
* supports {@link PeekingIterator#peek()}
*/
public static PeekingIterator peeking(final Iterator delegate) {
if (delegate instanceof PeekingIterator) {
return (PeekingIterator) delegate;
}
return new PeekingIterator<>(delegate);
}
public static Iterator filter(
final Iterator delegate,
final Predicate filter
) {
return new FilterIterator<>(delegate, filter);
}
public static KeyValueIterator filterKv(
final KeyValueIterator delegate,
final Predicate filter
) {
return new KvFilterIterator<>(delegate, filter);
}
/**
* Transforms the delegate iterator into a {@link KeyValueIterator}
* by applying the transformation method {@code extract(T)} passed
* in.
*/
public static KeyValueIterator kv(
final Iterator delegate,
final Function> extract
) {
return new TransformIterator<>(delegate, extract);
}
/**
* Wraps a {@link KeyValueIterator} that already returns what the
* {@link WindowStoreIterator} requires but using {@link WindowedKey}
* instead of {@code Long}.
*/
public static WindowStoreIterator windowed(
final KeyValueIterator delegate
) {
return new WindowIterator(delegate);
}
/**
* Returns an iterator that contains window ends instead of just
* the stamped start timestamp.
*/
public static KeyValueIterator, byte[]> windowedKey(
final KeyValueIterator delegate,
final long windowSize
) {
return new WindowKeyIterator(delegate, windowSize);
}
/**
* Returns an iterator that iterates over all delegates in order
*/
public static , V> KeyValueIterator wrapped(
final List> delegates
) {
return new MultiPartitionRangeIterator<>(delegates);
}
public static KeyValueIterator mapKeys(
final KeyValueIterator delegate,
final Function mapper
) {
return new KeyMappingIterator<>(delegate, mapper);
}
private static class PeekingIterator implements Iterator {
private final Iterator delegate;
private T cached;
private PeekingIterator(final Iterator delegate) {
this.delegate = delegate;
}
public T peek() {
cache();
return cached;
}
@Override
public boolean hasNext() {
cache();
return cached != null;
}
@Override
public T next() {
cache();
final T next = cached;
cached = null;
return next;
}
private void cache() {
if (cached == null && delegate.hasNext()) {
cached = delegate.next();
}
}
}
private static class TransformIterator implements KeyValueIterator {
private final PeekingIterator delegate;
private final Function> extract;
private TransformIterator(
final Iterator delegate,
final Function> extract
) {
if (delegate instanceof KeyValueIterator) {
throw new IllegalArgumentException("TransformIterator should not wrap KeyValueIterators");
}
this.delegate = peeking(delegate);
this.extract = extract;
}
@Override
public void close() {
}
@Override
public K peekNextKey() {
return extract.apply(delegate.peek()).key;
}
@Override
public boolean hasNext() {
return delegate.hasNext();
}
@Override
public KeyValue next() {
return extract.apply(delegate.next());
}
}
private static class WindowIterator implements WindowStoreIterator {
private final KeyValueIterator delegate;
public WindowIterator(final KeyValueIterator delegate) {
this.delegate = delegate;
}
@Override
public void close() {
delegate.close();
}
@Override
public Long peekNextKey() {
return delegate.peekNextKey().windowStartMs;
}
@Override
public boolean hasNext() {
return delegate.hasNext();
}
@Override
public KeyValue next() {
final KeyValue next = delegate.next();
return new KeyValue<>(next.key.windowStartMs, next.value);
}
}
private static class KvFilterIterator implements KeyValueIterator {
private final KeyValueIterator delegate;
private final Predicate filter;
private KvFilterIterator(
final KeyValueIterator delegate,
final Predicate filter
) {
this.delegate = delegate;
this.filter = filter;
}
@Override
public void close() {
delegate.close();
}
@Override
public K peekNextKey() {
if (hasNext()) {
return delegate.peekNextKey();
}
return null;
}
@Override
public boolean hasNext() {
while (delegate.hasNext()) {
if (filter.test(delegate.peekNextKey())) {
return true;
}
// otherwise filter it out and try again
delegate.next();
}
return false;
}
@Override
public KeyValue next() {
if (hasNext()) {
return delegate.next();
}
throw new NoSuchElementException();
}
}
private static class FilterIterator implements Iterator {
private final PeekingIterator delegate;
private final Predicate filter;
private FilterIterator(final Iterator delegate, final Predicate filter) {
this.delegate = peeking(delegate);
this.filter = filter;
}
@Override
public boolean hasNext() {
while (delegate.hasNext()) {
if (filter.test(delegate.peek())) {
return true;
}
// otherwise filter it out and try again
delegate.next();
}
return false;
}
@Override
public T next() {
if (hasNext()) {
return delegate.next();
}
throw new NoSuchElementException();
}
}
private static class KeyMappingIterator implements KeyValueIterator {
private final KeyValueIterator delegate;
private final Function mapper;
private KeyMappingIterator(final KeyValueIterator delegate, final Function mapper) {
this.delegate = delegate;
this.mapper = mapper;
}
@Override
public boolean hasNext() {
return delegate.hasNext();
}
@Override
public KeyValue next() {
if (hasNext()) {
final var next = delegate.next();
return new KeyValue<>(mapper.apply(next.key), next.value);
}
throw new NoSuchElementException();
}
@Override
public void close() {
delegate.close();
}
@Override
public O peekNextKey() {
final var next = delegate.peekNextKey();
return mapper.apply(next);
}
}
private static class WindowKeyIterator implements KeyValueIterator, byte[]> {
private final KeyValueIterator delegate;
private final long windowSize;
private WindowKeyIterator(
final KeyValueIterator delegate,
final long windowSize
) {
this.delegate = delegate;
this.windowSize = windowSize;
}
@Override
public void close() {
delegate.close();
}
@Override
public Windowed peekNextKey() {
return fromStamp(delegate.peekNextKey());
}
@Override
public boolean hasNext() {
return delegate.hasNext();
}
@Override
public KeyValue, byte[]> next() {
final KeyValue next = delegate.next();
return new KeyValue<>(
fromStamp(next.key),
next.value
);
}
private Windowed fromStamp(final WindowedKey windowedKey) {
return new Windowed<>(
windowedKey.key,
new TimeWindow(windowedKey.windowStartMs, endTs(windowedKey.windowStartMs))
);
}
private long endTs(final long stamp) {
final long endMs = stamp + windowSize;
return endMs < 0 ? Long.MAX_VALUE : endMs;
}
}
} © 2015 - 2024 Weber Informatics LLC | Privacy Policy