org.apache.samza.storage.kv.KeyValueIterator Java Examples

private void fanOut(String sender, Map<String, Object> message, MessageCollector collector) {
  // Colon is used as separator, and semicolon is lexicographically after colon
  KeyValueIterator<String, String> followers = socialGraph.range(sender + ":", sender + ";");

  try {
    while (followers.hasNext()) {
      String[] follow = followers.next().getKey().split(":");
      if (!follow[0].equals(sender)) {
        throw new IllegalStateException("Social graph db prefix doesn't match: " + sender + " != " + follow[0]);
      }
      message.put("recipient", follow[1]);
      collector.send(new OutgoingMessageEnvelope(NewsfeedConfig.DELIVERIES_STREAM, follow[1], null, message));
    }
  } finally {
    followers.close();
  }
}
 

@Override
public void window(MessageCollector collector,
  TaskCoordinator coordinator) {

  KeyValueIterator<String, String> entries = store.all();
  while (entries.hasNext()) {                                        // c
    Entry<String, String> entry = entries.next();
    String key = entry.getKey();
    String value = entry.getValue();
    if (isTimestampKey(key) && Cart.isAbandoned(value)) {            // d
      String shopper = extractShopper(key);
      String cart = store.get(asCartKey(shopper));
      
      AbandonedCartEvent event =
        new AbandonedCartEvent(shopper, cart);
      collector.send(new OutgoingMessageEnvelope(
        new SystemStream("kafka", "derived-events-ch04"), event));    // e
      
      resetShopper(shopper);
    }
  }
}
 

@SuppressWarnings("unchecked")
public void init(Context context) {
  this.store = (KeyValueStore<String, String>) context.getTaskContext().getStore("mystore");
  System.out.println("Contents of store: ");
  KeyValueIterator<String, String> iter = store.all();
  while (iter.hasNext()) {
    Entry<String, String> entry = iter.next();
    System.out.println(entry.getKey() + " => " + entry.getValue());
  }
  iter.close();
}
 

@Override
public void remove(K key, long startTimestamp, long endTimeStamp) {
  validateRange(startTimestamp, endTimeStamp);
  TimeSeriesKey<K> fromKey = new TimeSeriesKey(key, startTimestamp, 0);
  TimeSeriesKey<K> toKey = new TimeSeriesKey(key, endTimeStamp, 0);

  List<TimeSeriesKey<K>> keysToDelete = new LinkedList<>();

  KeyValueIterator<TimeSeriesKey<K>, V> range = kvStore.range(fromKey, toKey);
  try {
    while (range.hasNext()) {
      keysToDelete.add(range.next().getKey());
    }
  } finally {
    range.close();
  }
  kvStore.deleteAll(keysToDelete);
}
 

@Override
public void clear() {
  final ByteArray maxKey = createMaxKey();
  final KeyValueIterator<ByteArray, byte[]> kvIter = store.range(getEncodedStoreKey(), maxKey);
  while (kvIter.hasNext()) {
    store.delete(kvIter.next().getKey());
  }
  kvIter.close();
}
 

@Test
public void testSnapshot() throws Exception {
  InMemoryKeyValueStore store = new InMemoryKeyValueStore(
      new KeyValueStoreMetrics("testInMemory", new MetricsRegistryMap()));
  ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
  String prefix = "prefix";
  for (int i = 0; i < 100; i++) {
    store.put(genKey(outputStream, prefix, i), genValue());
  }

  byte[] firstKey = genKey(outputStream, prefix, 0);
  byte[] lastKey = genKey(outputStream, prefix, 100);
  KeyValueSnapshot<byte[], byte[]> snapshot = store.snapshot(firstKey, lastKey);
  // Make sure the cached Iterable won't change when new elements are added
  store.put(genKey(outputStream, prefix, 200), genValue());
  assertTrue(Iterators.size(snapshot.iterator()) == 100);

  List<Integer> keys = new ArrayList<>();
  KeyValueIterator<byte[], byte[]> iter = snapshot.iterator();
  while (iter.hasNext()) {
    Entry<byte[], byte[]> entry = iter.next();
    int key = Ints.fromByteArray(Arrays.copyOfRange(entry.getKey(), prefix.getBytes().length, entry.getKey().length));
    keys.add(key);
  }
  assertEquals(keys, IntStream.rangeClosed(0, 99).boxed().collect(Collectors.toList()));

  outputStream.close();
  store.close();
}
 

@Override
public void window(MessageCollector collector, TaskCoordinator coordinator) {
  String currentEpoch = this.store.get(CURRENT_EPOCH);
  logger.info("Checking if epoch " + currentEpoch + " is complete.");
  int count = 0;
  KeyValueIterator<String, String> iter = this.store.all();

  while (iter.hasNext()) {
    Entry<String, String> entry = iter.next();
    String foundEpoch = entry.getValue();
    if (foundEpoch.equals(currentEpoch)) {
      count += 1;
    } else {
      logger.info("####### Found a different epoch! - " + foundEpoch + " Current epoch is " + currentEpoch);
    }
  }
  iter.close();
  if (count == expectedKeys + 1) {
    logger.info("Epoch " + currentEpoch + " is complete.");
    int nextEpoch = Integer.parseInt(currentEpoch) + 1;
    for (int i = 0; i < numPartitions; i++) {
      logger.info("Emitting next epoch - " + Integer.toString(i) + " -> " + Integer.toString(nextEpoch));
      collector.send(new OutgoingMessageEnvelope(new SystemStream("kafka", "epoch"), Integer.toString(i), Integer.toString(nextEpoch)));
    }
    this.store.put(CURRENT_EPOCH, Integer.toString(nextEpoch));
  } else if (count > expectedKeys + 1) {
    throw new IllegalStateException("Got " + count + " keys, which is more than the expected " + (expectedKeys + 1));
  } else {
    logger.info("Only found " + count + " valid keys, try again later.");
  }
}
 

@Override
public KeyValueSnapshot<K, V> snapshot(K from, K to) {
  final ConcurrentNavigableMap<byte[], byte[]> values = map.subMap(keySerde.toBytes(from), keySerde.toBytes(to));
  return new KeyValueSnapshot<K, V>() {
    @Override
    public KeyValueIterator<K, V> iterator() {
      return new InMemoryIterator<>(values.entrySet().iterator(), keySerde, valSerde);
    }

    @Override
    public void close() {
    }
  };
}
 

@Override
public ClosableIterator<TimestampedValue<V>> get(K key, long startTimestamp, long endTimestamp) {
  validateRange(startTimestamp, endTimestamp);
  TimeSeriesKey<K> fromKey = new TimeSeriesKey(key, startTimestamp, 0);
  TimeSeriesKey<K> toKey = new TimeSeriesKey(key, endTimestamp, 0);

  KeyValueIterator<TimeSeriesKey<K>, V> range = kvStore.range(fromKey, toKey);

  LOG.trace("Getting entries in the store for {} from {} to {}", new Object[] {key, startTimestamp, endTimestamp});
  return new TimeSeriesStoreIterator<>(range);
}
 

@Override
public KeyValueIterator<K, V> all() {
  final Iterator<Map.Entry<K, V>> iterator = map.entrySet().iterator();
  return new KeyValueIterator<K, V>() {
    @Override
    public void close() {
      //not applicable
    }

    @Override
    public boolean hasNext() {
      return iterator.hasNext();
    }

    @Override
    public Entry<K, V> next() {
      Map.Entry<K, V> kv = iterator.next();
      return new Entry<>(kv.getKey(), kv.getValue());
    }

    @Override
    public void remove() {
      iterator.remove();
    }

  };
}
 

/**
 * Since we are not able to track the instances of the iterators created here and close them
 * properly, we need to load the content into memory.
 */
private <OutputT> Iterable<OutputT> createIterable(
    SerializableFunction<org.apache.samza.storage.kv.Entry<ByteArray, byte[]>, OutputT> fn) {
  final ByteArray maxKey = createMaxKey();
  final KeyValueIterator<ByteArray, byte[]> kvIter = store.range(getEncodedStoreKey(), maxKey);
  final List<Entry<ByteArray, byte[]>> iterable = ImmutableList.copyOf(kvIter);
  kvIter.close();

  return new Iterable<OutputT>() {
    @Override
    public Iterator<OutputT> iterator() {
      final Iterator<Entry<ByteArray, byte[]>> iter = iterable.iterator();

      return new Iterator<OutputT>() {
        @Override
        public boolean hasNext() {
          return iter.hasNext();
        }

        @Override
        public OutputT next() {
          return fn.apply(iter.next());
        }
      };
    }
  };
}
 

public TimeSeriesStoreIterator(KeyValueIterator<TimeSeriesKey<K>, V> wrappedIterator) {
  this.wrappedIterator = wrappedIterator;
}
 

@Override
public KeyValueIterator<K, V> range(K from, K to) {
  ConcurrentNavigableMap<byte[], byte[]> values = map.subMap(keySerde.toBytes(from), keySerde.toBytes(to));
  return new InMemoryIterator(values.entrySet().iterator(), keySerde, valSerde);
}
 

@Override
public KeyValueIterator<K, V> all() {
  return new InMemoryIterator(map.entrySet().iterator(), keySerde, valSerde);
}
 

@Override
public KeyValueIterator<K, V> range(K from, K to) {
  throw new RuntimeException("not implemented.");
}
 

TestKeyValueIteraor(KeyValueIterator<byte[], byte[]> iter) {
  this.iter = iter;
}
 

@Override
public KeyValueIterator<byte[], byte[]> range(byte[] from, byte[] to) {
  TestKeyValueIteraor iter = new TestKeyValueIteraor(super.range(from, to));
  iterators.add(iter);
  return iter;
}
 

@Override
public void closeIterators() {
  openIterators.forEach(KeyValueIterator::close);
  openIterators.clear();
}