org.apache.beam.fn.harness.MapFnRunners Maven / Gradle / Ivy
/*
* 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.fn.harness;
import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables.getOnlyElement;
import java.io.IOException;
import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform;
import org.apache.beam.sdk.fn.data.FnDataReceiver;
import org.apache.beam.sdk.function.ThrowingFunction;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Iterables;
/**
* Utilities to create {@code PTransformRunners} which execute simple map functions.
*
* Simple map functions are used in a large number of transforms, especially runner-managed
* transforms, such as map_windows.
*
*
TODO: Add support for DoFns which are actually user supplied map/lambda functions instead of
* using the {@link FnApiDoFnRunner} instance.
*/
@SuppressWarnings({
"rawtypes" // TODO(https://github.com/apache/beam/issues/20447)
})
public abstract class MapFnRunners {
/** Create a {@link MapFnRunners} where the map function consumes elements directly. */
public static PTransformRunnerFactory forValueMapFnFactory(
ValueMapFnFactory fnFactory) {
return new Factory<>(new CompressedValueOnlyMapperFactory<>(fnFactory));
}
/**
* Create a {@link MapFnRunners} where the map function consumes {@link WindowedValue Windowed
* Values} and produced {@link WindowedValue Windowed Values}.
*
* Each {@link WindowedValue} provided to the function produced by the {@link
* WindowedValueMapFnFactory} will be in exactly one {@link
* org.apache.beam.sdk.transforms.windowing.BoundedWindow window}.
*/
public static PTransformRunnerFactory forWindowedValueMapFnFactory(
WindowedValueMapFnFactory fnFactory) {
return new Factory<>(new ExplodedWindowedValueMapperFactory<>(fnFactory));
}
/** A function factory which given a PTransform returns a map function. */
public interface ValueMapFnFactory {
ThrowingFunction forPTransform(String ptransformId, PTransform pTransform)
throws IOException;
}
/**
* A function factory which given a PTransform returns a map function over the entire {@link
* WindowedValue} of input and output elements.
*
* {@link WindowedValue Windowed Values} will only ever be in a single window.
*/
public interface WindowedValueMapFnFactory {
ThrowingFunction, WindowedValue> forPTransform(
String ptransformId, PTransform ptransform) throws IOException;
}
/** A factory for {@link MapFnRunners}s. */
private static class Factory
implements PTransformRunnerFactory> {
private final MapperFactory mapperFactory;
private Factory(MapperFactory mapperFactory) {
this.mapperFactory = mapperFactory;
}
@Override
public Mapper createRunnerForPTransform(Context context) throws IOException {
FnDataReceiver> consumer =
context.getPCollectionConsumer(
getOnlyElement(context.getPTransform().getOutputsMap().values()));
Mapper mapper =
mapperFactory.create(context.getPTransformId(), context.getPTransform(), consumer);
String pCollectionId =
Iterables.getOnlyElement(context.getPTransform().getInputsMap().values());
context.addPCollectionConsumer(pCollectionId, mapper::map);
return mapper;
}
}
@FunctionalInterface
private interface MapperFactory {
Mapper create(
String ptransformId, PTransform ptransform, FnDataReceiver> outputs)
throws IOException;
}
private interface Mapper {
void map(WindowedValue input) throws Exception;
}
private static class ExplodedWindowedValueMapperFactory
implements MapperFactory {
private final WindowedValueMapFnFactory fnFactory;
private ExplodedWindowedValueMapperFactory(
WindowedValueMapFnFactory fnFactory) {
this.fnFactory = fnFactory;
}
@Override
public Mapper create(
String ptransformId, PTransform ptransform, FnDataReceiver> outputs)
throws IOException {
ThrowingFunction, WindowedValue> fn =
fnFactory.forPTransform(ptransformId, ptransform);
return input -> {
for (WindowedValue exploded : input.explodeWindows()) {
outputs.accept(fn.apply(exploded));
}
};
}
}
private static class CompressedValueOnlyMapperFactory
implements MapperFactory {
private final ValueMapFnFactory fnFactory;
private CompressedValueOnlyMapperFactory(ValueMapFnFactory fnFactory) {
this.fnFactory = fnFactory;
}
@Override
public Mapper create(
String ptransformId, PTransform ptransform, FnDataReceiver> outputs)
throws IOException {
ThrowingFunction fn = fnFactory.forPTransform(ptransformId, ptransform);
return input -> outputs.accept(input.withValue(fn.apply(input.getValue())));
}
}
}