@Override public <T> void output(TupleTag<T> tag, WindowedValue<T> value) { if (!openBuffer) { emit(tag, value); } else { bufferState.add(KV.<Integer, WindowedValue<?>>of(tagsToIds.get(tag), value)); } }
@Override public <T> void output(TupleTag<T> tag, WindowedValue<T> value) { if (!openBuffer) { emit(tag, value); } else { bufferState.add(KV.of(tagsToIds.get(tag), value)); } }
@Override public <T> void output(TupleTag<T> tag, WindowedValue<T> value) { if (!openBuffer) { emit(tag, value); } else { bufferState.add(KV.of(tagsToIds.get(tag), value)); } }
@Override public void append(K key, W window, Iterator<V> values) { initStateInternals(key); StateNamespace namespace = StateNamespaces.window(windowCoder, window); BagState<V> bagState = stateInternals.state(namespace, stateTag); while (values.hasNext()) { bagState.add(values.next()); } }
@Override public void append(K key, W window, Iterator<V> values) { initStateInternals(key); StateNamespace namespace = StateNamespaces.window(windowCoder, window); BagState<V> bagState = stateInternals.state(namespace, stateTag); while (values.hasNext()) { bagState.add(values.next()); } }
@Test public void testMergeBagIntoSource() throws Exception { BagState<String> bag1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR); BagState<String> bag2 = underTest.state(NAMESPACE_2, STRING_BAG_ADDR); bag1.add("Hello"); bag2.add("World"); bag1.add("!"); StateMerging.mergeBags(Arrays.asList(bag1, bag2), bag1); // Reading the merged bag gets both the contents assertThat(bag1.read(), containsInAnyOrder("Hello", "World", "!")); assertThat(bag2.read(), Matchers.emptyIterable()); }
@Test public void testBagWithBadCoderEquality() throws Exception { // Ensure two instances of the bad coder are distinct; models user who fails to // override equals() or inherit from CustomCoder for StructuredCoder assertThat( new StringCoderWithIdentityEquality(), not(equalTo(new StringCoderWithIdentityEquality()))); BagState<String> state1 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR1); state1.add("hello"); BagState<String> state2 = underTest.state(NAMESPACE_1, STRING_BAG_ADDR2); assertThat(state2.read(), containsInAnyOrder("hello")); }
@Override public final void processElement1( StreamRecord<WindowedValue<InputT>> streamRecord) throws Exception { checkInvokeStartBundle(); Iterable<WindowedValue<InputT>> justPushedBack = pushbackDoFnRunner.processElementInReadyWindows(streamRecord.getValue()); BagState<WindowedValue<InputT>> pushedBack = nonKeyedStateInternals.state(StateNamespaces.global(), pushedBackTag); checkInitPushedBackWatermark(); long min = pushedBackWatermark.get(); for (WindowedValue<InputT> pushedBackValue : justPushedBack) { min = Math.min(min, pushedBackValue.getTimestamp().getMillis()); pushedBack.add(pushedBackValue); } setPushedBackWatermark(min); checkInvokeFinishBundleByCount(); }
@Test public void testBag() throws Exception { BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR); // State instances are cached, but depend on the namespace. assertThat(value, equalTo(underTest.state(NAMESPACE_1, STRING_BAG_ADDR))); assertThat(value, not(equalTo(underTest.state(NAMESPACE_2, STRING_BAG_ADDR)))); assertThat(value.read(), Matchers.emptyIterable()); value.add("hello"); assertThat(value.read(), containsInAnyOrder("hello")); value.add("world"); assertThat(value.read(), containsInAnyOrder("hello", "world")); value.clear(); assertThat(value.read(), Matchers.emptyIterable()); assertThat(underTest.state(NAMESPACE_1, STRING_BAG_ADDR), equalTo(value)); }
@ProcessElement public void processElement( @Element KV<String, Integer> element, @StateId(stateId) BagState<MyInteger> state, OutputReceiver<List<MyInteger>> r) { state.add(new MyInteger(element.getValue())); Iterable<MyInteger> currentValue = state.read(); if (Iterables.size(currentValue) >= 4) { List<MyInteger> sorted = Lists.newArrayList(currentValue); Collections.sort(sorted); r.output(sorted); } } };
@Test public void testBagIsEmpty() throws Exception { BagState<String> value = underTest.state(NAMESPACE_1, STRING_BAG_ADDR); assertThat(value.isEmpty().read(), Matchers.is(true)); ReadableState<Boolean> readFuture = value.isEmpty(); value.add("hello"); assertThat(readFuture.read(), Matchers.is(false)); value.clear(); assertThat(readFuture.read(), Matchers.is(true)); }
@ProcessElement public void processElement( @Element KV<String, Integer> element, @StateId(stateId) BagState<MyInteger> state, OutputReceiver<List<MyInteger>> r) { state.add(new MyInteger(element.getValue())); Iterable<MyInteger> currentValue = state.read(); if (Iterables.size(currentValue) >= 4) { List<MyInteger> sorted = Lists.newArrayList(currentValue); Collections.sort(sorted); r.output(sorted); } } };
@Test public void testGetWithEmpty() { CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null); StateNamespace namespace = new StateNamespaceForTest("foo"); StateTag<BagState<String>> bagTag = StateTags.bag("foo", StringUtf8Coder.of()); BagState<String> stringBag = internals.state(namespace, bagTag); assertThat(stringBag.read(), emptyIterable()); stringBag.add("bar"); stringBag.add("baz"); assertThat(stringBag.read(), containsInAnyOrder("baz", "bar")); BagState<String> reReadStringBag = internals.state(namespace, bagTag); assertThat(reReadStringBag.read(), containsInAnyOrder("baz", "bar")); }
@Test public void testCommitWithEmptyNewAndFullUnderlyingIsNotEmpty() { CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null); CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying); StateNamespace namespace = new StateNamespaceForTest("foo"); StateTag<BagState<String>> bagTag = StateTags.bag("foo", StringUtf8Coder.of()); BagState<String> stringBag = underlying.state(namespace, bagTag); assertThat(stringBag.read(), emptyIterable()); stringBag.add("bar"); stringBag.add("baz"); internals.commit(); assertThat(internals.isEmpty(), is(false)); }
@Test public void testCommitWithEmptyNewAndFullUnderlyingIsNotEmpty() { CopyOnAccessInMemoryStateInternals<String> underlying = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null); CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, underlying); StateNamespace namespace = new StateNamespaceForTest("foo"); StateTag<BagState<String>> bagTag = StateTags.bag("foo", StringUtf8Coder.of()); BagState<String> stringBag = underlying.state(namespace, bagTag); assertThat(stringBag.read(), emptyIterable()); stringBag.add("bar"); stringBag.add("baz"); internals.commit(); assertThat(internals.isEmpty(), is(false)); }
@Test public void testCommitWithOnlyClearedValuesIsEmpty() { CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null); StateNamespace namespace = new StateNamespaceForTest("foo"); StateTag<BagState<String>> bagTag = StateTags.bag("foo", StringUtf8Coder.of()); BagState<String> stringBag = internals.state(namespace, bagTag); assertThat(stringBag.read(), emptyIterable()); stringBag.add("foo"); stringBag.clear(); internals.commit(); assertThat(internals.isEmpty(), is(true)); }
@Test public void testCommitWithOnlyClearedValuesIsEmpty() { CopyOnAccessInMemoryStateInternals<String> internals = CopyOnAccessInMemoryStateInternals.withUnderlying(key, null); StateNamespace namespace = new StateNamespaceForTest("foo"); StateTag<BagState<String>> bagTag = StateTags.bag("foo", StringUtf8Coder.of()); BagState<String> stringBag = internals.state(namespace, bagTag); assertThat(stringBag.read(), emptyIterable()); stringBag.add("foo"); stringBag.clear(); internals.commit(); assertThat(internals.isEmpty(), is(true)); }
@Test public void getExecutionContextDifferentKeysIndependentState() { StepStateAndTimers fooContext = context.getStateAndTimers(createdProducer, StructuralKey.of("foo", StringUtf8Coder.of())); StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of()); fooContext.stateInternals().state(StateNamespaces.global(), intBag).add(1); StepStateAndTimers barContext = context.getStateAndTimers(createdProducer, StructuralKey.of("bar", StringUtf8Coder.of())); assertThat(barContext, not(equalTo(fooContext))); assertThat( barContext.stateInternals().state(StateNamespaces.global(), intBag).read(), emptyIterable()); }
@Test public void getExecutionContextDifferentStepsIndependentState() { StructuralKey<?> myKey = StructuralKey.of("foo", StringUtf8Coder.of()); StepStateAndTimers fooContext = context.getStateAndTimers(createdProducer, myKey); StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of()); fooContext.stateInternals().state(StateNamespaces.global(), intBag).add(1); StepStateAndTimers barContext = context.getStateAndTimers(downstreamProducer, myKey); assertThat( barContext.stateInternals().state(StateNamespaces.global(), intBag).read(), emptyIterable()); }
@Test public void getExecutionContextDifferentStepsIndependentState() { StructuralKey<?> myKey = StructuralKey.of("foo", StringUtf8Coder.of()); DirectExecutionContext fooContext = context.getExecutionContext(createdProducer, myKey); StateTag<BagState<Integer>> intBag = StateTags.bag("myBag", VarIntCoder.of()); fooContext.getStepContext("s1").stateInternals().state(StateNamespaces.global(), intBag).add(1); DirectExecutionContext barContext = context.getExecutionContext(downstreamProducer, myKey); assertThat( barContext .getStepContext("s1") .stateInternals() .state(StateNamespaces.global(), intBag) .read(), emptyIterable()); }