/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.beam.runners.core;
import com.google.common.base.MoreObjects;
import java.io.IOException;
import java.io.Serializable;
import java.util.Objects;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderRegistry;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.CombiningState;
import org.apache.beam.sdk.state.MapState;
import org.apache.beam.sdk.state.SetState;
import org.apache.beam.sdk.state.State;
import org.apache.beam.sdk.state.StateBinder;
import org.apache.beam.sdk.state.StateSpec;
import org.apache.beam.sdk.state.StateSpecs;
import org.apache.beam.sdk.state.ValueState;
import org.apache.beam.sdk.state.WatermarkHoldState;
import org.apache.beam.sdk.transforms.Combine.CombineFn;
import org.apache.beam.sdk.transforms.CombineWithContext.CombineFnWithContext;
import org.apache.beam.sdk.transforms.windowing.BoundedWindow;
import org.apache.beam.sdk.transforms.windowing.TimestampCombiner;
/**
* Static utility methods for creating {@link StateTag} instances.
*/
@Experimental(Kind.STATE)
public class StateTags {
private static final CoderRegistry STANDARD_REGISTRY = CoderRegistry.createDefault();
/** @deprecated for migration purposes only */
@Deprecated
private static StateBinder adaptTagBinder(final StateTag.StateBinder binder) {
return new StateBinder() {
@Override
public <T> ValueState<T> bindValue(
String id, StateSpec<ValueState<T>> spec, Coder<T> coder) {
return binder.bindValue(tagForSpec(id, spec), coder);
}
@Override
public <T> BagState<T> bindBag(
String id, StateSpec<BagState<T>> spec, Coder<T> elemCoder) {
return binder.bindBag(tagForSpec(id, spec), elemCoder);
}
@Override
public <T> SetState<T> bindSet(
String id, StateSpec<SetState<T>> spec, Coder<T> elemCoder) {
return binder.bindSet(tagForSpec(id, spec), elemCoder);
}
@Override
public <KeyT, ValueT> MapState<KeyT, ValueT> bindMap(
String id, StateSpec<MapState<KeyT, ValueT>> spec,
Coder<KeyT> mapKeyCoder, Coder<ValueT> mapValueCoder) {
return binder.bindMap(tagForSpec(id, spec), mapKeyCoder, mapValueCoder);
}
@Override
public <InputT, AccumT, OutputT>
CombiningState<InputT, AccumT, OutputT> bindCombining(
String id,
StateSpec<CombiningState<InputT, AccumT, OutputT>> spec,
Coder<AccumT> accumCoder,
CombineFn<InputT, AccumT, OutputT> combineFn) {
return binder.bindCombiningValue(tagForSpec(id, spec), accumCoder, combineFn);
}
@Override
public <InputT, AccumT, OutputT>
CombiningState<InputT, AccumT, OutputT> bindCombiningWithContext(
String id,
StateSpec<CombiningState<InputT, AccumT, OutputT>> spec,
Coder<AccumT> accumCoder,
CombineFnWithContext<InputT, AccumT, OutputT> combineFn) {
return binder.bindCombiningValueWithContext(
tagForSpec(id, spec), accumCoder, combineFn);
}
@Override
public WatermarkHoldState bindWatermark(
String id,
StateSpec<WatermarkHoldState> spec,
TimestampCombiner timestampCombiner) {
return binder.bindWatermark(tagForSpec(id, spec), timestampCombiner);
}
};
}
private enum StateKind {
SYSTEM('s'),
USER('u');
private char prefix;
StateKind(char prefix) {
this.prefix = prefix;
}
}
private StateTags() { }
private interface SystemStateTag<StateT extends State> {
StateTag<StateT> asKind(StateKind kind);
}
/** Create a state tag for the given id and spec. */
public static <StateT extends State> StateTag<StateT> tagForSpec(
String id, StateSpec<StateT> spec) {
return new SimpleStateTag<>(new StructuredId(id), spec);
}
/**
* Create a simple state tag for values of type {@code T}.
*/
public static <T> StateTag<ValueState<T>> value(String id, Coder<T> valueCoder) {
return new SimpleStateTag<>(new StructuredId(id), StateSpecs.value(valueCoder));
}
/**
* Create a state tag for values that use a {@link CombineFn} to automatically merge
* multiple {@code InputT}s into a single {@code OutputT}.
*/
public static <InputT, AccumT, OutputT>
StateTag<CombiningState<InputT, AccumT, OutputT>>
combiningValue(
String id, Coder<AccumT> accumCoder, CombineFn<InputT, AccumT, OutputT> combineFn) {
return new SimpleStateTag<>(
new StructuredId(id), StateSpecs.combining(accumCoder, combineFn));
}
/**
* Create a state tag for values that use a {@link CombineFnWithContext} to automatically
* merge multiple {@code InputT}s into a single {@code OutputT}.
*/
public static <InputT, AccumT, OutputT>
StateTag<CombiningState<InputT, AccumT, OutputT>>
combiningValueWithContext(
String id,
Coder<AccumT> accumCoder,
CombineFnWithContext<InputT, AccumT, OutputT> combineFn) {
return new SimpleStateTag<>(
new StructuredId(id), StateSpecs.combining(accumCoder, combineFn));
}
/**
* Create a state tag for values that use a {@link CombineFn} to automatically merge
* multiple {@code InputT}s into a single {@code OutputT}.
*
* <p>This determines the {@code Coder<AccumT>} from the given {@code Coder<InputT>}, and
* should only be used to initialize static values.
*/
public static <InputT, AccumT, OutputT>
StateTag<CombiningState<InputT, AccumT, OutputT>>
combiningValueFromInputInternal(
String id, Coder<InputT> inputCoder, CombineFn<InputT, AccumT, OutputT> combineFn) {
return new SimpleStateTag<>(
new StructuredId(id), StateSpecs.combiningFromInputInternal(inputCoder, combineFn));
}
/**
* Create a state tag that is optimized for adding values frequently, and
* occasionally retrieving all the values that have been added.
*/
public static <T> StateTag<BagState<T>> bag(String id, Coder<T> elemCoder) {
return new SimpleStateTag<>(new StructuredId(id), StateSpecs.bag(elemCoder));
}
/**
* Create a state spec that supporting for {@link java.util.Set} like access patterns.
*/
public static <T> StateTag<SetState<T>> set(String id, Coder<T> elemCoder) {
return new SimpleStateTag<>(new StructuredId(id), StateSpecs.set(elemCoder));
}
/**
* Create a state spec that supporting for {@link java.util.Map} like access patterns.
*/
public static <K, V> StateTag<MapState<K, V>> map(
String id, Coder<K> keyCoder, Coder<V> valueCoder) {
return new SimpleStateTag<>(new StructuredId(id), StateSpecs.map(keyCoder, valueCoder));
}
/**
* Create a state tag for holding the watermark.
*/
public static <W extends BoundedWindow> StateTag<WatermarkHoldState>
watermarkStateInternal(String id, TimestampCombiner timestampCombiner) {
return new SimpleStateTag<>(
new StructuredId(id), StateSpecs.watermarkStateInternal(timestampCombiner));
}
/**
* Convert an arbitrary {@link StateTag} to a system-internal tag that is guaranteed not to
* collide with any user tags.
*/
public static <StateT extends State> StateTag<StateT> makeSystemTagInternal(
StateTag<StateT> tag) {
if (!(tag instanceof SystemStateTag)) {
throw new IllegalArgumentException("Expected subclass of SimpleStateTag, got " + tag);
}
// Checked above
@SuppressWarnings("unchecked")
SystemStateTag<StateT> typedTag = (SystemStateTag<StateT>) tag;
return typedTag.asKind(StateKind.SYSTEM);
}
public static <InputT, AccumT, OutputT> StateTag<BagState<AccumT>>
convertToBagTagInternal(
StateTag<CombiningState<InputT, AccumT, OutputT>> combiningTag) {
return new SimpleStateTag<>(
new StructuredId(combiningTag.getId()),
StateSpecs.convertToBagSpecInternal(combiningTag.getSpec()));
}
private static class StructuredId implements Serializable {
private final StateKind kind;
private final String rawId;
private StructuredId(String rawId) {
this(StateKind.USER, rawId);
}
private StructuredId(StateKind kind, String rawId) {
this.kind = kind;
this.rawId = rawId;
}
public StructuredId asKind(StateKind kind) {
return new StructuredId(kind, rawId);
}
public void appendTo(Appendable sb) throws IOException {
sb.append(kind.prefix).append(rawId);
}
public String getRawId() {
return rawId;
}
@Override
public String toString() {
return MoreObjects.toStringHelper(getClass())
.add("id", rawId)
.add("kind", kind)
.toString();
}
@Override
public boolean equals(Object obj) {
if (obj == this) {
return true;
}
if (!(obj instanceof StructuredId)) {
return false;
}
StructuredId that = (StructuredId) obj;
return Objects.equals(this.kind, that.kind)
&& Objects.equals(this.rawId, that.rawId);
}
@Override
public int hashCode() {
return Objects.hash(kind, rawId);
}
}
/**
* A basic {@link StateTag} implementation that manages the structured ids.
*/
private static class SimpleStateTag<StateT extends State>
implements StateTag<StateT>, SystemStateTag<StateT> {
private final StateSpec<StateT> spec;
private final StructuredId id;
public SimpleStateTag(StructuredId id, StateSpec<StateT> spec) {
this.id = id;
this.spec = spec;
}
@Override
@Deprecated
public StateT bind(StateTag.StateBinder binder) {
return spec.bind(
this.id.getRawId(), adaptTagBinder(binder));
}
@Override
public String getId() {
return id.getRawId();
}
@Override
public StateSpec<StateT> getSpec() {
return spec;
}
@Override
public String toString() {
return MoreObjects.toStringHelper(getClass())
.add("id", id)
.toString();
}
@Override
public void appendTo(Appendable sb) throws IOException {
id.appendTo(sb);
}
@Override
public StateTag<StateT> asKind(StateKind kind) {
return new SimpleStateTag<>(id.asKind(kind), spec);
}
@Override
public boolean equals(Object other) {
if (!(other instanceof SimpleStateTag)) {
return false;
}
SimpleStateTag<?> otherTag = (SimpleStateTag<?>) other;
return Objects.equals(this.getId(), otherTag.getId())
&& Objects.equals(this.getSpec(), otherTag.getSpec());
}
@Override
public int hashCode() {
return Objects.hash(getClass(), this.getId(), this.getSpec());
}
}
}