com.google.cloud.dataflow.sdk.util.AppliedCombineFn Maven / Gradle / Ivy
Show all versions of google-cloud-dataflow-java-sdk-all Show documentation
/*
* Copyright (C) 2015 Google Inc.
*
* Licensed 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 com.google.cloud.dataflow.sdk.util;
import com.google.cloud.dataflow.sdk.coders.CannotProvideCoderException;
import com.google.cloud.dataflow.sdk.coders.Coder;
import com.google.cloud.dataflow.sdk.coders.CoderRegistry;
import com.google.cloud.dataflow.sdk.coders.KvCoder;
import com.google.cloud.dataflow.sdk.transforms.CombineFnBase.PerKeyCombineFn;
import com.google.cloud.dataflow.sdk.transforms.CombineWithContext.KeyedCombineFnWithContext;
import com.google.cloud.dataflow.sdk.values.PCollectionView;
import com.google.common.annotations.VisibleForTesting;
import java.io.Serializable;
/**
* A {@link KeyedCombineFnWithContext} with a fixed accumulator coder. This is created from a
* specific application of the {@link KeyedCombineFnWithContext}.
*
* Because the {@code AccumT} may reference {@code InputT}, the specific {@code Coder}
* may depend on the {@code Coder}.
*
* @param type of keys
* @param type of input values
* @param type of mutable accumulator values
* @param type of output values
*/
public class AppliedCombineFn implements Serializable {
private final PerKeyCombineFn fn;
private final Coder accumulatorCoder;
private final Iterable> sideInputViews;
private final KvCoder kvCoder;
private final WindowingStrategy windowingStrategy;
private AppliedCombineFn(PerKeyCombineFn fn,
Coder accumulatorCoder, Iterable> sideInputViews,
KvCoder kvCoder, WindowingStrategy windowingStrategy) {
this.fn = fn;
this.accumulatorCoder = accumulatorCoder;
this.sideInputViews = sideInputViews;
this.kvCoder = kvCoder;
this.windowingStrategy = windowingStrategy;
}
public static AppliedCombineFn
withAccumulatorCoder(
PerKeyCombineFn fn,
Coder accumCoder) {
return withAccumulatorCoder(fn, accumCoder, null, null, null);
}
public static AppliedCombineFn
withAccumulatorCoder(
PerKeyCombineFn fn,
Coder accumCoder, Iterable> sideInputViews,
KvCoder kvCoder, WindowingStrategy windowingStrategy) {
// Casting down the K and InputT is safe because they're only used as inputs.
@SuppressWarnings("unchecked")
PerKeyCombineFn clonedFn =
(PerKeyCombineFn) SerializableUtils.clone(fn);
return create(clonedFn, accumCoder, sideInputViews, kvCoder, windowingStrategy);
}
@VisibleForTesting
public static AppliedCombineFn
withInputCoder(PerKeyCombineFn fn,
CoderRegistry registry, KvCoder kvCoder) {
return withInputCoder(fn, registry, kvCoder, null, null);
}
public static AppliedCombineFn
withInputCoder(PerKeyCombineFn fn,
CoderRegistry registry, KvCoder kvCoder,
Iterable> sideInputViews, WindowingStrategy windowingStrategy) {
// Casting down the K and InputT is safe because they're only used as inputs.
@SuppressWarnings("unchecked")
PerKeyCombineFn clonedFn =
(PerKeyCombineFn) SerializableUtils.clone(fn);
try {
Coder accumulatorCoder = clonedFn.getAccumulatorCoder(
registry, kvCoder.getKeyCoder(), kvCoder.getValueCoder());
return create(clonedFn, accumulatorCoder, sideInputViews, kvCoder, windowingStrategy);
} catch (CannotProvideCoderException e) {
throw new IllegalStateException("Could not determine coder for accumulator", e);
}
}
private static AppliedCombineFn create(
PerKeyCombineFn fn,
Coder accumulatorCoder, Iterable> sideInputViews,
KvCoder kvCoder, WindowingStrategy windowingStrategy) {
return new AppliedCombineFn<>(
fn, accumulatorCoder, sideInputViews, kvCoder, windowingStrategy);
}
public PerKeyCombineFn getFn() {
return fn;
}
public Iterable> getSideInputViews() {
return sideInputViews;
}
public Coder getAccumulatorCoder() {
return accumulatorCoder;
}
public KvCoder getKvCoder() {
return kvCoder;
}
public WindowingStrategy getWindowingStrategy() {
return windowingStrategy;
}
}