diff --git a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto index 9360522ab409..fde3657b667d 100644 --- a/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto +++ b/model/fn-execution/src/main/proto/org/apache/beam/model/fn_execution/v1/beam_fn_api.proto @@ -1009,6 +1009,29 @@ message StateKey { bytes key = 4; } + // Represents a request for all of the entries of a multimap associated with a + // specified user key and window for a PTransform. See + // https://s.apache.org/beam-fn-state-api-and-bundle-processing for further + // details. + // + // Can only be used to perform StateGetRequests and StateClearRequests on the + // user state. + // + // The response data stream will be a concatenation of pairs, where the first + // component is the map key and the second component is a concatenation of + // values associated with that map key. + message MultimapEntriesUserState { + // (Required) The id of the PTransform containing user state. + string transform_id = 1; + // (Required) The id of the user state. + string user_state_id = 2; + // (Required) The window encoded in a nested context. + bytes window = 3; + // (Required) The key of the currently executing element encoded in a + // nested context. + bytes key = 4; + } + // Represents a request for the values of the map key associated with a // specified user key and window for a PTransform. See // https://s.apache.org/beam-fn-state-api-and-bundle-processing for further @@ -1064,6 +1087,7 @@ message StateKey { MultimapKeysSideInput multimap_keys_side_input = 5; MultimapKeysValuesSideInput multimap_keys_values_side_input = 8; MultimapKeysUserState multimap_keys_user_state = 6; + MultimapEntriesUserState multimap_entries_user_state = 10; MultimapUserState multimap_user_state = 7; OrderedListUserState ordered_list_user_state = 9; } diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java index 8409133772eb..e3464d98526c 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java @@ -2917,6 +2917,72 @@ public void processElement( pipeline.run(); } + @Test + @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMultimapState.class}) + public void testMultimapStateEntries() { + final String stateId = "foo:"; + final String countStateId = "count"; + DoFn>, KV> fn = + new DoFn>, KV>() { + + @StateId(stateId) + private final StateSpec> multimapState = + StateSpecs.multimap(StringUtf8Coder.of(), VarIntCoder.of()); + + @StateId(countStateId) + private final StateSpec> countState = + StateSpecs.combiningFromInputInternal(VarIntCoder.of(), Sum.ofIntegers()); + + @ProcessElement + public void processElement( + ProcessContext c, + @Element KV> element, + @StateId(stateId) MultimapState state, + @StateId(countStateId) CombiningState count, + OutputReceiver> r) { + // Empty before we process any elements. + if (count.read() == 0) { + assertThat(state.entries().read(), emptyIterable()); + } + assertEquals(count.read().intValue(), Iterables.size(state.entries().read())); + + KV value = element.getValue(); + state.put(value.getKey(), value.getValue()); + count.add(1); + + if (count.read() >= 4) { + // Those should be evaluated only when ReadableState.read is called. + ReadableState>> entriesView = state.entries(); + ReadableState containsBadView = state.containsKey("BadKey"); + + // Those are evaluated immediately. + Iterable> entries = state.entries().read(); + + state.remove("b"); + assertEquals(4, Iterables.size(entries)); + state.put("a", 2); + + // Note we output the view of state before the modifications in this if statement. + for (Entry entry : entries) { + r.output(KV.of(entry.getKey(), entry.getValue())); + } + } + } + }; + PCollection> output = + pipeline + .apply( + Create.of( + KV.of("hello", KV.of("a", 97)), KV.of("hello", KV.of("a", 97)), + KV.of("hello", KV.of("a", 98)), KV.of("hello", KV.of("b", 33)))) + .apply(ParDo.of(fn)); + PAssert.that(output) + .containsInAnyOrder( + KV.of("a", 97), KV.of("a", 97), + KV.of("a", 98), KV.of("b", 33)); + pipeline.run(); + } + @Test @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesMultimapState.class}) public void testMultimapStateRemove() { diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java index e06a82c8e25f..6913c75a5f2d 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/FnApiStateAccessor.java @@ -117,7 +117,7 @@ public static class Factory { public Factory( PipelineOptions pipelineOptions, - Set runnerCapabilites, + Set runnerCapabilities, String ptransformId, Supplier processBundleInstructionId, Supplier> cacheTokens, @@ -128,7 +128,7 @@ public Factory( Coder keyCoder, Coder windowCoder) { this.pipelineOptions = pipelineOptions; - this.runnerCapabilities = runnerCapabilites; + this.runnerCapabilities = runnerCapabilities; this.ptransformId = ptransformId; this.processBundleInstructionId = processBundleInstructionId; this.cacheTokens = cacheTokens; @@ -240,7 +240,7 @@ public FnApiStateAccessor create() { } private final PipelineOptions pipelineOptions; - private final Set runnerCapabilites; + private final Set runnerCapabilities; private final Map stateKeyObjectCache; private final Map, SideInputSpec> sideInputSpecMap; private final BeamFnStateClient beamFnStateClient; @@ -259,7 +259,7 @@ public FnApiStateAccessor create() { public FnApiStateAccessor( PipelineOptions pipelineOptions, - Set runnerCapabilites, + Set runnerCapabilities, String ptransformId, Supplier processBundleInstructionId, Supplier> cacheTokens, @@ -270,7 +270,7 @@ public FnApiStateAccessor( Coder keyCoder, Coder windowCoder) { this.pipelineOptions = pipelineOptions; - this.runnerCapabilites = runnerCapabilites; + this.runnerCapabilities = runnerCapabilities; this.stateKeyObjectCache = Maps.newHashMap(); this.sideInputSpecMap = sideInputSpecMap; this.beamFnStateClient = beamFnStateClient; @@ -414,7 +414,7 @@ public T get(PCollectionView view, BoundedWindow window) { key, ((KvCoder) sideInputSpec.getCoder()).getKeyCoder(), ((KvCoder) sideInputSpec.getCoder()).getValueCoder(), - runnerCapabilites.contains( + runnerCapabilities.contains( BeamUrns.getUrn( RunnerApi.StandardRunnerProtocols.Enum .MULTIMAP_KEYS_VALUES_SIDE_INPUT)))); @@ -762,8 +762,113 @@ public MultimapState bindMultimap( StateSpec> spec, Coder keyCoder, Coder valueCoder) { - // TODO(https://github.com/apache/beam/issues/23616) - throw new UnsupportedOperationException("Multimap is not currently supported with Fn API."); + return (MultimapState) + stateKeyObjectCache.computeIfAbsent( + createMultimapKeysUserStateKey(id), + new Function() { + @Override + public Object apply(StateKey stateKey) { + return new MultimapState() { + private final MultimapUserState impl = + createMultimapUserState(stateKey, keyCoder, valueCoder); + + @Override + public void put(KeyT key, ValueT value) { + impl.put(key, value); + } + + @Override + public ReadableState> get(KeyT key) { + return new ReadableState>() { + @Override + public Iterable read() { + return impl.get(key); + } + + @Override + public ReadableState> readLater() { + impl.get(key).prefetch(); + return this; + } + }; + } + + @Override + public void remove(KeyT key) { + impl.remove(key); + } + + @Override + public ReadableState> keys() { + return new ReadableState>() { + @Override + public Iterable read() { + return impl.keys(); + } + + @Override + public ReadableState> readLater() { + impl.keys().prefetch(); + return this; + } + }; + } + + @Override + public ReadableState>> entries() { + return new ReadableState>>() { + @Override + public Iterable> read() { + return impl.entries(); + } + + @Override + public ReadableState>> readLater() { + impl.entries().prefetch(); + return this; + } + }; + } + + @Override + public ReadableState containsKey(KeyT key) { + return new ReadableState() { + @Override + public Boolean read() { + return !Iterables.isEmpty(impl.get(key)); + } + + @Override + public ReadableState readLater() { + impl.get(key).prefetch(); + return this; + } + }; + } + + @Override + public ReadableState isEmpty() { + return new ReadableState() { + @Override + public Boolean read() { + return Iterables.isEmpty(impl.keys()); + } + + @Override + public ReadableState readLater() { + impl.keys().prefetch(); + return this; + } + }; + } + + @Override + public void clear() { + impl.clear(); + } + }; + } + }); } @Override diff --git a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java index 617faba87cc0..967cb76af1f3 100644 --- a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java +++ b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/state/MultimapUserState.java @@ -38,12 +38,15 @@ import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateRequest; import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.IterableCoder; +import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.fn.stream.PrefetchableIterable; import org.apache.beam.sdk.fn.stream.PrefetchableIterables; import org.apache.beam.sdk.fn.stream.PrefetchableIterator; import org.apache.beam.sdk.util.ByteStringOutputStream; import org.apache.beam.sdk.values.KV; import org.apache.beam.vendor.grpc.v1p69p0.com.google.protobuf.ByteString; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; /** @@ -52,9 +55,6 @@ * *

Calling {@link #asyncClose()} schedules any required persistence changes. This object should * no longer be used after it is closed. - * - *

TODO: Move to an async persist model where persistence is signalled based upon cache memory - * pressure and its need to flush. */ public class MultimapUserState { @@ -63,8 +63,10 @@ public class MultimapUserState { private final Coder mapKeyCoder; private final Coder valueCoder; private final StateRequest keysStateRequest; + private final StateRequest entriesStateRequest; private final StateRequest userStateRequest; private final CachingStateIterable persistedKeys; + private final CachingStateIterable>> persistedEntries; private boolean isClosed; private boolean isCleared; @@ -106,6 +108,23 @@ public MultimapUserState( .setWindow(stateKey.getMultimapKeysUserState().getWindow()) .setKey(stateKey.getMultimapKeysUserState().getKey()); this.userStateRequest = userStateRequestBuilder.build(); + + StateRequest.Builder entriesStateRequestBuilder = StateRequest.newBuilder(); + entriesStateRequestBuilder + .setInstructionId(instructionId) + .getStateKeyBuilder() + .getMultimapEntriesUserStateBuilder() + .setTransformId(stateKey.getMultimapKeysUserState().getTransformId()) + .setUserStateId(stateKey.getMultimapKeysUserState().getUserStateId()) + .setWindow(stateKey.getMultimapKeysUserState().getWindow()) + .setKey(stateKey.getMultimapKeysUserState().getKey()); + this.entriesStateRequest = entriesStateRequestBuilder.build(); + this.persistedEntries = + StateFetchingIterators.readAllAndDecodeStartingFrom( + Caches.subCache(this.cache, "AllEntries"), + beamFnStateClient, + entriesStateRequest, + KvCoder.of(mapKeyCoder, IterableCoder.of(valueCoder))); } public void clear() { @@ -235,6 +254,122 @@ public K next() { }; } + @SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/21068) + }) + /* + * Returns an Iterable containing all entries in this multimap. + */ + public PrefetchableIterable> entries() { + checkState( + !isClosed, + "Multimap user state is no longer usable because it is closed for %s", + keysStateRequest.getStateKey()); + if (isCleared) { + Map>> pendingAddsNow = new HashMap<>(pendingAdds); + return PrefetchableIterables.concat( + Iterables.concat( + Iterables.transform( + pendingAddsNow.entrySet(), + entry -> + Iterables.transform( + entry.getValue().getValue(), + value -> Maps.immutableEntry(entry.getValue().getKey(), value))))); + } + + Set pendingRemovesNow = new HashSet<>(pendingRemoves.keySet()); + // Make a deep copy of pendingAdds so this iterator represents a snapshot of state at the time + // it was created. + Map> pendingAddsNow = new HashMap<>(); + for (Map.Entry>> entry : pendingAdds.entrySet()) { + pendingAddsNow.put( + entry.getValue().getKey(), new ArrayList<>()); // entry.getValue().getValue()); + for (V value : entry.getValue().getValue()) { + List values = pendingAddsNow.get(entry.getValue().getKey()); + values.add(value); + } + } + return new PrefetchableIterables.Default>() { + @Override + public PrefetchableIterator> createIterator() { + return new PrefetchableIterator>() { + PrefetchableIterator> persistedEntriesIterator = + PrefetchableIterables.concat( + Iterables.concat( + Iterables.transform( + persistedEntries, + entry -> + Iterables.transform( + entry.getValue(), + value -> Maps.immutableEntry(entry.getKey(), value))))) + .iterator(); + Iterator> pendingAddsNowIterator; + boolean hasNext; + Map.Entry nextEntry; + + @Override + public boolean isReady() { + return persistedEntriesIterator.isReady(); + } + + @Override + public void prefetch() { + if (!isReady()) { + persistedEntriesIterator.prefetch(); + } + } + + @Override + public boolean hasNext() { + if (hasNext) { + return true; + } + + while (persistedEntriesIterator.hasNext()) { + Map.Entry nextPersistedEntry = persistedEntriesIterator.next(); + Object nextKeyStructuralValue = + mapKeyCoder.structuralValue(nextPersistedEntry.getKey()); + if (!pendingRemovesNow.contains(nextKeyStructuralValue)) { + nextEntry = + Maps.immutableEntry(nextPersistedEntry.getKey(), nextPersistedEntry.getValue()); + hasNext = true; + return true; + } + } + + if (pendingAddsNowIterator == null) { + pendingAddsNowIterator = + Iterables.concat( + Iterables.transform( + pendingAddsNow.entrySet(), + entry -> + Iterables.transform( + entry.getValue(), + value -> Maps.immutableEntry(entry.getKey(), value)))) + .iterator(); + } + while (pendingAddsNowIterator.hasNext()) { + nextEntry = pendingAddsNowIterator.next(); + hasNext = true; + return true; + } + + return false; + } + + @Override + public Map.Entry next() { + if (!hasNext()) { + throw new NoSuchElementException(); + } + hasNext = false; + return nextEntry; + } + }; + } + }; + } + /* * Store a key-value pair in the multimap. * Allows duplicate key-value pairs. diff --git a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java index 48c9ce43bdf0..457362815cf6 100644 --- a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java +++ b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/state/MultimapUserStateTest.java @@ -22,6 +22,7 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.emptyIterable; import static org.hamcrest.collection.ArrayMatching.arrayContainingInAnyOrder; +import static org.hamcrest.collection.IsIterableContainingInAnyOrder.containsInAnyOrder; import static org.hamcrest.core.Is.is; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; @@ -34,11 +35,15 @@ import java.util.Collections; import java.util.Iterator; import java.util.Map; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; import org.apache.beam.fn.harness.Cache; import org.apache.beam.fn.harness.Caches; import org.apache.beam.model.fnexecution.v1.BeamFnApi.StateKey; import org.apache.beam.sdk.coders.ByteArrayCoder; import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.IterableCoder; +import org.apache.beam.sdk.coders.KvCoder; import org.apache.beam.sdk.coders.NullableCoder; import org.apache.beam.sdk.coders.StringUtf8Coder; import org.apache.beam.sdk.fn.stream.PrefetchableIterable; @@ -179,6 +184,66 @@ public void testKeys() throws Exception { assertThrows(IllegalStateException.class, () -> userState.keys()); } + @Test + public void testEntries() throws Exception { + FakeBeamFnStateClient fakeClient = + new FakeBeamFnStateClient( + ImmutableMap.of( + createMultimapEntriesStateKey(), + KV.of( + KvCoder.of(ByteArrayCoder.of(), IterableCoder.of(StringUtf8Coder.of())), + asList(KV.of(A1, asList("V1", "V2")), KV.of(A2, asList("V3")))))); + MultimapUserState userState = + new MultimapUserState<>( + Caches.noop(), + fakeClient, + "instructionId", + createMultimapKeyStateKey(), + ByteArrayCoder.of(), + StringUtf8Coder.of()); + + assertArrayEquals(A1, userState.entries().iterator().next().getKey()); + assertThat( + StreamSupport.stream(userState.entries().spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder( + KV.of(ByteString.copyFrom(A1), "V1"), + KV.of(ByteString.copyFrom(A1), "V2"), + KV.of(ByteString.copyFrom(A2), "V3"))); + + userState.put(A1, "V4"); + assertThat( + StreamSupport.stream(userState.entries().spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder( + KV.of(ByteString.copyFrom(A1), "V1"), + KV.of(ByteString.copyFrom(A1), "V2"), + KV.of(ByteString.copyFrom(A2), "V3"), + KV.of(ByteString.copyFrom(A1), "V4"))); + + userState.remove(A1); + assertThat( + StreamSupport.stream(userState.entries().spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder(KV.of(ByteString.copyFrom(A2), "V3"))); + + userState.put(A1, "V5"); + assertThat( + StreamSupport.stream(userState.entries().spliterator(), false) + .map(entry -> KV.of(ByteString.copyFrom(entry.getKey()), entry.getValue())) + .collect(Collectors.toList()), + containsInAnyOrder( + KV.of(ByteString.copyFrom(A2), "V3"), KV.of(ByteString.copyFrom(A1), "V5"))); + + userState.clear(); + assertThat(userState.entries(), emptyIterable()); + userState.asyncClose(); + assertThrows(IllegalStateException.class, () -> userState.entries()); + } + @Test public void testPut() throws Exception { FakeBeamFnStateClient fakeClient = @@ -1053,6 +1118,17 @@ private StateKey createMultimapKeyStateKey() throws IOException { .build(); } + private StateKey createMultimapEntriesStateKey() throws IOException { + return StateKey.newBuilder() + .setMultimapEntriesUserState( + StateKey.MultimapEntriesUserState.newBuilder() + .setWindow(encode(encodedWindow)) + .setKey(encode(encodedKey)) + .setTransformId(pTransformId) + .setUserStateId(stateId)) + .build(); + } + private StateKey createMultimapValueStateKey(byte[] key) throws IOException { return StateKey.newBuilder() .setMultimapUserState( diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/DetectNewPartitionsAction.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/DetectNewPartitionsAction.java index 080372d04593..d340a9ee4aad 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/DetectNewPartitionsAction.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/spanner/changestreams/action/DetectNewPartitionsAction.java @@ -189,14 +189,17 @@ private void outputBatch( final PartitionMetadata updatedPartition = partition.toBuilder().setScheduledAt(scheduledAt).build(); + Instant outputTimestamp = new Instant(minWatermark.toSqlTimestamp()); LOG.info( - "[{}] Outputting partition at {} with start time {} and end time {}", + "[{}] Outputting partition at {} with start time {}, end time {}, creation time {} and output timestamp {}", updatedPartition.getPartitionToken(), updatedPartition.getScheduledAt(), updatedPartition.getStartTimestamp(), - updatedPartition.getEndTimestamp()); + updatedPartition.getEndTimestamp(), + createdAt, + outputTimestamp); - receiver.outputWithTimestamp(partition, new Instant(minWatermark.toSqlTimestamp())); + receiver.outputWithTimestamp(partition, outputTimestamp); metrics.incPartitionRecordCount(); metrics.updatePartitionCreatedToScheduled(