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

static Factory createStateInternalFactory(
    String id,
    Coder<?> keyCoder,
    TaskContext context,
    SamzaPipelineOptions pipelineOptions,
    DoFnSignature signature) {
  final int batchGetSize = pipelineOptions.getStoreBatchGetSize();
  final Map<String, KeyValueStore<ByteArray, byte[]>> stores = new HashMap<>();
  stores.put(BEAM_STORE, getBeamStore(context));

  final Coder stateKeyCoder;
  if (keyCoder != null) {
    signature
        .stateDeclarations()
        .keySet()
        .forEach(
            stateId ->
                stores.put(
                    stateId, (KeyValueStore<ByteArray, byte[]>) context.getStore(stateId)));
    stateKeyCoder = keyCoder;
  } else {
    stateKeyCoder = VoidCoder.of();
  }
  return new Factory<>(Objects.toString(id), stores, stateKeyCoder, batchGetSize);
}
 

protected AbstractSamzaState(
    StateNamespace namespace, StateTag<? extends State> address, Coder<T> coder) {
  this.coder = coder;
  this.namespace = namespace.stringKey();

  final KeyValueStore<ByteArray, byte[]> userStore = stores.get(address.getId());
  this.store = userStore != null ? userStore : stores.get(BEAM_STORE);

  final ByteArrayOutputStream baos = getThreadLocalBaos();
  try (DataOutputStream dos = new DataOutputStream(baos)) {
    dos.write(keyBytes);
    dos.writeUTF(namespace.stringKey());

    if (userStore == null) {
      // for system state, we need to differentiate based on the following:
      dos.writeUTF(stageId);
      dos.writeUTF(address.getId());
    }
  } catch (IOException e) {
    throw new RuntimeException(
        "Could not encode full address for state: " + address.getId(), e);
  }
  this.encodedStoreKey = baos.toByteArray();
}
 

/**
 * Processes the incoming side input message envelope and updates the last processed offset for its SSP.
 * Synchronized inorder to be exclusive with flush().
 *
 * @param envelope incoming envelope to be processed
 */
public synchronized void process(IncomingMessageEnvelope envelope) {
  SystemStreamPartition envelopeSSP = envelope.getSystemStreamPartition();
  String envelopeOffset = envelope.getOffset();

  for (String store: this.sspToStores.get(envelopeSSP)) {
    SideInputsProcessor storeProcessor = this.storeToProcessor.get(store);
    KeyValueStore keyValueStore = (KeyValueStore) this.taskSideInputStorageManager.getStore(store);
    Collection<Entry<?, ?>> entriesToBeWritten = storeProcessor.process(envelope, keyValueStore);

    // TODO: SAMZA-2255: optimize writes to side input stores
    for (Entry entry : entriesToBeWritten) {
      // If the key is null we ignore, if the value is null, we issue a delete, else we issue a put
      if (entry.getKey() != null) {
        if (entry.getValue() != null) {
          keyValueStore.put(entry.getKey(), entry.getValue());
        } else {
          keyValueStore.delete(entry.getKey());
        }
      }
    }
  }

  this.lastProcessedOffsets.put(envelopeSSP, envelopeOffset);
}
 

private static KeyValueStore<ByteArray, byte[]> createStore(String name) {
  final Options options = new Options();
  options.setCreateIfMissing(true);

  RocksDbKeyValueStore rocksStore =
      new RocksDbKeyValueStore(
          new File(System.getProperty("java.io.tmpdir") + "/" + name),
          options,
          new MapConfig(),
          false,
          "beamStore",
          new WriteOptions(),
          new FlushOptions(),
          new KeyValueStoreMetrics("beamStore", new MetricsRegistryMap()));

  return new SerializedKeyValueStore<>(
      rocksStore,
      new ByteArraySerdeFactory.ByteArraySerde(),
      new ByteSerde(),
      new SerializedKeyValueStoreMetrics("beamStore", new MetricsRegistryMap()));
}
 

private static SamzaTimerInternalsFactory<String> createTimerInternalsFactory(
    Scheduler<KeyedTimerData<String>> timerRegistry,
    String timerStateId,
    SamzaPipelineOptions pipelineOptions,
    KeyValueStore<ByteArray, byte[]> store) {

  final SamzaStoreStateInternals.Factory<?> nonKeyedStateInternalsFactory =
      createNonKeyedStateInternalsFactory(pipelineOptions, store);

  return SamzaTimerInternalsFactory.createTimerInternalFactory(
      StringUtf8Coder.of(),
      timerRegistry,
      timerStateId,
      nonKeyedStateInternalsFactory,
      (WindowingStrategy) WindowingStrategy.globalDefault(),
      PCollection.IsBounded.BOUNDED,
      pipelineOptions);
}
 

@Override
protected void handleInit(Context context) {

  KeyValueStore<TimeSeriesKey<K>, Object> store =
      (KeyValueStore<TimeSeriesKey<K>, Object>) context.getTaskContext().getStore(windowOpSpec.getOpId());

  if (initializer != null) {
    initializer.init(context);
  }

  if (keyFn != null) {
    keyFn.init(context);
  }

  // For aggregating windows, we use the store in over-write mode since we only retain the aggregated
  // value. Else, we use the store in append-mode.
  if (foldLeftFn != null) {
    foldLeftFn.init(context);
    timeSeriesStore = new TimeSeriesStoreImpl(store, false);
  } else {
    timeSeriesStore = new TimeSeriesStoreImpl(store, true);
  }
}
 

@Override
protected CompletionStage<Collection<JM>> handleMessageAsync(M message, MessageCollector collector,
    TaskCoordinator coordinator) {
  Collection<JM> output = Collections.emptyList();

  try {
    KeyValueStore<K, TimestampedValue<M>> thisState = thisPartialJoinFn.getState();
    KeyValueStore<K, TimestampedValue<OM>> otherState = otherPartialJoinFn.getState();

    K key = thisPartialJoinFn.getKey(message);
    thisState.put(key, new TimestampedValue<>(message, clock.currentTimeMillis()));
    TimestampedValue<OM> otherMessage = otherState.get(key);

    long now = clock.currentTimeMillis();
    if (otherMessage != null && otherMessage.getTimestamp() > now - ttlMs) {
      JM joinResult = thisPartialJoinFn.apply(message, otherMessage.getValue());
      output = Collections.singletonList(joinResult);
    }
  } catch (Exception e) {
    throw new SamzaException("Error handling message in PartialJoinOperatorImpl " + getOpImplId(), e);
  }

  return CompletableFuture.completedFuture(output);
}
 

@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();
}
 

public TaskContextImpl(TaskModel taskModel,
    MetricsRegistry taskMetricsRegistry,
    Function<String, KeyValueStore> keyValueStoreProvider,
    TableManager tableManager,
    CallbackScheduler callbackScheduler,
    OffsetManager offsetManager,
    JobModel jobModel,
    StreamMetadataCache streamMetadataCache) {
  this.taskModel = taskModel;
  this.taskMetricsRegistry = taskMetricsRegistry;
  this.keyValueStoreProvider = keyValueStoreProvider;
  this.tableManager = tableManager;
  this.callbackScheduler = callbackScheduler;
  this.offsetManager = offsetManager;
  this.jobModel = jobModel;
  this.streamMetadataCache = streamMetadataCache;
}
 

private SamzaStoreStateInternals(
    Map<String, KeyValueStore<ByteArray, byte[]>> stores,
    @Nullable K key,
    @Nullable byte[] keyBytes,
    String stageId,
    int batchGetSize) {
  this.stores = stores;
  this.key = key;
  this.keyBytes = keyBytes;
  this.batchGetSize = batchGetSize;
  this.stageId = stageId;
}
 

/**
 * Given that there is not a store corresponding to the storeName, getStore should throw an exception.
 */
@Test(expected = IllegalArgumentException.class)
public void testGetMissingStore() {
  KeyValueStore store = mock(KeyValueStore.class);
  when(keyValueStoreProvider.apply("myStore")).thenReturn(null);
  assertEquals(store, taskContext.getStore("myStore"));
}
 

/**
 * Given that there is a store corresponding to the storeName, getStore should return the store.
 */
@Test
public void testGetStore() {
  KeyValueStore store = mock(KeyValueStore.class);
  when(keyValueStoreProvider.apply("myStore")).thenReturn(store);
  assertEquals(store, taskContext.getStore("myStore"));
}
 

@Override
@SuppressWarnings("unchecked")
public void init(Context context) {
  this.store = (KeyValueStore<String, String>) context.getTaskContext().getStore("checker-state");
  this.expectedKeys = context.getJobContext().getConfig().getInt("expected.keys");
  this.numPartitions = context.getJobContext().getConfig().getInt("num.partitions");
}
 

private PartialJoinFunction<Object, Object, Object, Object> createRightJoinFn(JoinOperatorSpec joinOpSpec) {
  return new PartialJoinFunction<Object, Object, Object, Object>() {
    private final JoinFunction joinFn = joinOpSpec.getJoinFn();
    private KeyValueStore<Object, TimestampedValue<Object>> rightStreamState;

    @Override
    public Object apply(Object m, Object om) {
      return joinFn.apply(om, m);
    }

    @Override
    public Object getKey(Object message) {
      return joinFn.getSecondKey(message);
    }

    @Override
    public void init(Context context) {
      String rightStoreName = joinOpSpec.getRightOpId();
      rightStreamState =
          (KeyValueStore<Object, TimestampedValue<Object>>) context.getTaskContext().getStore(rightStoreName);

      // user-defined joinFn should only be initialized once,
      // so we do it only in left partial join function and not here again.
    }

    @Override
    public KeyValueStore<Object, TimestampedValue<Object>> getState() {
      return rightStreamState;
    }
  };
}
 

private PartialJoinFunction<Object, Object, Object, Object> createLeftJoinFn(JoinOperatorSpec joinOpSpec) {
  return new PartialJoinFunction<Object, Object, Object, Object>() {
    private final JoinFunction joinFn = joinOpSpec.getJoinFn();
    private KeyValueStore<Object, TimestampedValue<Object>> leftStreamState;

    @Override
    public Object apply(Object m, Object om) {
      return joinFn.apply(m, om);
    }

    @Override
    public Object getKey(Object message) {
      return joinFn.getFirstKey(message);
    }

    @Override
    public KeyValueStore<Object, TimestampedValue<Object>> getState() {
      return leftStreamState;
    }

    @Override
    public void init(Context context) {
      String leftStoreName = joinOpSpec.getLeftOpId();
      leftStreamState =
          (KeyValueStore<Object, TimestampedValue<Object>>) context.getTaskContext().getStore(leftStoreName);

      // user-defined joinFn should only be initialized once, so we do it only in left partial join function.
      joinFn.init(context);
    }

    @Override
    public void close() {
      // joinFn#close() must only be called once, so we do it it only in left partial join function.
      joinFn.close();
    }
  };
}
 

@Override
public KeyValueStore getStore(String storeName) {
  KeyValueStore store = this.keyValueStoreProvider.apply(storeName);
  if (store == null) {
    throw new IllegalArgumentException(String.format("No store found for storeName: %s", storeName));
  }
  return store;
}
 

/**
 * {@inheritDoc}
 * Override {@link org.apache.samza.operators.functions.InitableFunction#init(Context)} to
 * get a KeyValueStore for persistence and the MetricsRegistry for metrics.
 */
@Override
public void init(Context context) {
  TaskContext taskContext = context.getTaskContext();
  store = (KeyValueStore<String, Integer>) taskContext.getStore("wikipedia-stats");
  repeatEdits = taskContext.getTaskMetricsRegistry().newCounter("edit-counters", "repeat-edits");
}
 

@Test
public void testEventTimeTimers() {
  final SamzaPipelineOptions pipelineOptions =
      PipelineOptionsFactory.create().as(SamzaPipelineOptions.class);

  final KeyValueStore<ByteArray, byte[]> store = createStore("store1");
  final SamzaTimerInternalsFactory<String> timerInternalsFactory =
      createTimerInternalsFactory(null, "timer", pipelineOptions, store);

  final StateNamespace nameSpace = StateNamespaces.global();
  final TimerInternals timerInternals = timerInternalsFactory.timerInternalsForKey("testKey");
  final TimerInternals.TimerData timer1 =
      TimerInternals.TimerData.of(
          "timer1", nameSpace, new Instant(10), new Instant(10), TimeDomain.EVENT_TIME);
  timerInternals.setTimer(timer1);

  final TimerInternals.TimerData timer2 =
      TimerInternals.TimerData.of(
          "timer2", nameSpace, new Instant(100), new Instant(100), TimeDomain.EVENT_TIME);
  timerInternals.setTimer(timer2);

  timerInternalsFactory.setInputWatermark(new Instant(5));
  Collection<KeyedTimerData<String>> readyTimers = timerInternalsFactory.removeReadyTimers();
  assertTrue(readyTimers.isEmpty());

  timerInternalsFactory.setInputWatermark(new Instant(20));
  readyTimers = timerInternalsFactory.removeReadyTimers();
  assertEquals(1, readyTimers.size());
  assertEquals(timer1, readyTimers.iterator().next().getTimerData());

  timerInternalsFactory.setInputWatermark(new Instant(150));
  readyTimers = timerInternalsFactory.removeReadyTimers();
  assertEquals(1, readyTimers.size());
  assertEquals(timer2, readyTimers.iterator().next().getTimerData());

  store.close();
}
 

private static SamzaStoreStateInternals.Factory<?> createNonKeyedStateInternalsFactory(
    SamzaPipelineOptions pipelineOptions, KeyValueStore<ByteArray, byte[]> store) {
  final TaskContext context = mock(TaskContext.class);
  when(context.getStore(anyString())).thenReturn((KeyValueStore) store);
  final TupleTag<?> mainOutputTag = new TupleTag<>("output");

  return SamzaStoreStateInternals.createStateInternalFactory(
      "42", null, context, pipelineOptions, null);
}
 

@Override
public KeyValueStore<byte[], byte[]> getKVStore(
    String storeName,
    File storeDir,
    MetricsRegistry registry,
    SystemStreamPartition changeLogSystemStreamPartition,
    JobContext jobContext,
    ContainerContext containerContext,
    StorageEngineFactory.StoreMode readWrite) {
  KeyValueStoreMetrics metrics = new KeyValueStoreMetrics(storeName, registry);
  return new TestStore(metrics);
}
 

public Factory(
    String stageId,
    Map<String, KeyValueStore<ByteArray, byte[]>> stores,
    Coder<K> keyCoder,
    int batchGetSize) {
  this.stageId = stageId;
  this.stores = stores;
  this.keyCoder = keyCoder;
  this.batchGetSize = batchGetSize;
}
 

public void init(Config config, TaskContext context) {
  this.store = (KeyValueStore<String, String>)
    context.getStore("nile-carts");
}
 

public void init(Context context) {
  TaskContext taskContext = context.getTaskContext();
  this.store = (KeyValueStore<String, Integer>) taskContext.getStore("wikipedia-stats");
  this.repeatEdits = taskContext.getTaskMetricsRegistry().newCounter("edit-counters", "repeat-edits");
}
 

@Override
@SuppressWarnings("unchecked")
public void init(Config config, TaskContext context) throws Exception {
  homeTimeline = (KeyValueStore<String, Map<String, Object>>) context.getStore("home-timeline");
}
 

@Override
public void init(Context context) {
  this.store = (KeyValueStore<String, String>) context.getTaskContext().getStore("joiner-state");
  this.expected = context.getJobContext().getConfig().getInt("num.partitions");
  this.taskName = context.getTaskContext().getTaskModel().getTaskName();
}
 

@Override
public void init(Context context) {
  this.state = (KeyValueStore<String, String>) context.getTaskContext().getStore("emitter-state");
  this.taskName = context.getTaskContext().getTaskModel().getTaskName();
  this.max = context.getJobContext().getConfig().getInt("count");
}
 

@SuppressWarnings("unchecked")
public void init(Context context) {
  this.store = (KeyValueStore<String, String>) context.getTaskContext().getStore("mystore");
}
 

@Override
@SuppressWarnings("unchecked")
public void init(Config config, TaskContext context) throws Exception {
  socialGraph = (KeyValueStore<String, String>) context.getStore("social-graph");
  userTimeline = (KeyValueStore<String, Map<String, Object>>) context.getStore("user-timeline");
}
 

@Override
public void init(Context context) {
  this.statsStore =
      (KeyValueStore<String, StatsWindowState>) context.getTaskContext().getStore("my-stats-wnd-store");
}
 

private static <K> TimeSeriesStore<K, byte[]> newTimeSeriesStore(Serde<K> keySerde, boolean appendMode) {
  KeyValueStore<TimeSeriesKey<K>, byte[]> kvStore =
      new TestInMemoryStore(new TimeSeriesKeySerde<>(keySerde), new ByteSerde());
  return new TimeSeriesStoreImpl<>(kvStore, appendMode);
}