org.apache.flink.table.runtime.functions.ExecutionContextImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of flink-table_2.11
The newest version!
/*
 * 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.flink.table.runtime.functions;

import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.SortedMapStateDescriptor;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.common.typeutils.base.ListSerializer;
import org.apache.flink.api.common.typeutils.base.MapSerializer;
import org.apache.flink.api.common.typeutils.base.SortedMapSerializer;
import org.apache.flink.runtime.state.keyed.KeyedListState;
import org.apache.flink.runtime.state.keyed.KeyedListStateDescriptor;
import org.apache.flink.runtime.state.keyed.KeyedMapState;
import org.apache.flink.runtime.state.keyed.KeyedMapStateDescriptor;
import org.apache.flink.runtime.state.keyed.KeyedSortedMapState;
import org.apache.flink.runtime.state.keyed.KeyedSortedMapStateDescriptor;
import org.apache.flink.runtime.state.keyed.KeyedState;
import org.apache.flink.runtime.state.keyed.KeyedStateDescriptor;
import org.apache.flink.runtime.state.keyed.KeyedValueState;
import org.apache.flink.runtime.state.keyed.KeyedValueStateDescriptor;
import org.apache.flink.runtime.state.subkeyed.SubKeyedListState;
import org.apache.flink.runtime.state.subkeyed.SubKeyedListStateDescriptor;
import org.apache.flink.runtime.state.subkeyed.SubKeyedMapState;
import org.apache.flink.runtime.state.subkeyed.SubKeyedMapStateDescriptor;
import org.apache.flink.runtime.state.subkeyed.SubKeyedSortedMapState;
import org.apache.flink.runtime.state.subkeyed.SubKeyedSortedMapStateDescriptor;
import org.apache.flink.runtime.state.subkeyed.SubKeyedState;
import org.apache.flink.runtime.state.subkeyed.SubKeyedStateDescriptor;
import org.apache.flink.runtime.state.subkeyed.SubKeyedValueState;
import org.apache.flink.runtime.state.subkeyed.SubKeyedValueStateDescriptor;
import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
import org.apache.flink.table.dataformat.BaseRow;
import org.apache.flink.table.dataview.KeyedStateListView;
import org.apache.flink.table.dataview.KeyedStateMapView;
import org.apache.flink.table.dataview.KeyedStateSortedMapView;
import org.apache.flink.table.dataview.NullAwareKeyedStateMapView;
import org.apache.flink.table.dataview.NullAwareSubKeyedStateMapView;
import org.apache.flink.table.dataview.StateDataView;
import org.apache.flink.table.dataview.StateListView;
import org.apache.flink.table.dataview.StateMapView;
import org.apache.flink.table.dataview.StateSortedMapView;
import org.apache.flink.table.dataview.SubKeyedStateListView;
import org.apache.flink.table.dataview.SubKeyedStateMapView;
import org.apache.flink.table.typeutils.ListViewTypeInfo;
import org.apache.flink.table.typeutils.MapViewTypeInfo;
import org.apache.flink.table.typeutils.SortedMapViewTypeInfo;
import org.apache.flink.util.Preconditions;

import java.util.ArrayList;
import java.util.List;


/**
 * Implementation of ExecutionContext.
 */
@SuppressWarnings("unchecked")
public final class ExecutionContextImpl implements ExecutionContext {

	private static final String NULL_STATE_POSTFIX = "_null_state";

	private final AbstractStreamOperator operator;
	private final RuntimeContext runtimeContext;
	private final TypeSerializer namespaceSerializer;
	private final List> registeredStateDataViews;

	public ExecutionContextImpl(
			AbstractStreamOperator operator,
			RuntimeContext runtimeContext) {
		this(operator, runtimeContext, null);
	}

	public ExecutionContextImpl(
			AbstractStreamOperator operator,
			RuntimeContext runtimeContext,
			TypeSerializer namespaceSerializer) {
		this.operator = operator;
		this.runtimeContext = Preconditions.checkNotNull(runtimeContext);
		this.namespaceSerializer = namespaceSerializer;
		this.registeredStateDataViews = new ArrayList<>();
	}

	@Override
	public > S getKeyedState(KeyedStateDescriptor descriptor) throws Exception {
		return operator.getKeyedState(descriptor);
	}

	@Override
	public > S getSubKeyedState(SubKeyedStateDescriptor descriptor) throws Exception {
		return operator.getSubKeyedState(descriptor);
	}

	@Override
	public  KeyedValueState getKeyedValueState(
		ValueStateDescriptor descriptor) throws Exception {
		descriptor.initializeSerializerUnlessSet(operator.getExecutionConfig());
		return operator.getKeyedState(
			new KeyedValueStateDescriptor<>(
				descriptor.getName(),
				(TypeSerializer) operator.getKeySerializer(),
				descriptor.getSerializer()
			)
		);
	}

	@Override
	public  KeyedListState getKeyedListState(
		ListStateDescriptor descriptor
	) throws Exception {
		descriptor.initializeSerializerUnlessSet(operator.getExecutionConfig());
		return operator.getKeyedState(
			new KeyedListStateDescriptor<>(
				descriptor.getName(),
				(TypeSerializer) operator.getKeySerializer(),
				(ListSerializer) descriptor.getSerializer()
			)
		);
	}

	@Override
	public  KeyedMapState getKeyedMapState(
		MapStateDescriptor descriptor
	) throws Exception {
		descriptor.initializeSerializerUnlessSet(operator.getExecutionConfig());
		return operator.getKeyedState(
			new KeyedMapStateDescriptor<>(
				descriptor.getName(),
				(TypeSerializer) operator.getKeySerializer(),
				(MapSerializer) descriptor.getSerializer()
			)
		);
	}

	@Override
	public  KeyedSortedMapState getKeyedSortedMapState(
		SortedMapStateDescriptor descriptor
	) throws Exception {
		descriptor.initializeSerializerUnlessSet(operator.getExecutionConfig());
		return operator.getKeyedState(
			new KeyedSortedMapStateDescriptor<>(
				descriptor.getName(),
				(TypeSerializer) operator.getKeySerializer(),
				(SortedMapSerializer) descriptor.getSerializer()
			)
		);
	}

	@Override
	public  SubKeyedValueState getSubKeyedValueState(
		ValueStateDescriptor descriptor
	) throws Exception {
		if (namespaceSerializer == null) {
			throw new RuntimeException("The namespace serializer has not been initialized.");
		}

		descriptor.initializeSerializerUnlessSet(operator.getExecutionConfig());
		return operator.getSubKeyedState(
			new SubKeyedValueStateDescriptor<>(
				descriptor.getName(),
				(TypeSerializer) operator.getKeySerializer(),
				(TypeSerializer) namespaceSerializer,
				descriptor.getSerializer()
			)
		);
	}

	@Override
	public  SubKeyedListState getSubKeyedListState(
		ListStateDescriptor descriptor
	) throws Exception {
		if (namespaceSerializer == null) {
			throw new RuntimeException("The namespace serializer has not been initialized.");
		}

		descriptor.initializeSerializerUnlessSet(operator.getExecutionConfig());
		return operator.getSubKeyedState(new SubKeyedListStateDescriptor<>(
			descriptor.getName(),
			(TypeSerializer) operator.getKeySerializer(),
			(TypeSerializer) namespaceSerializer,
			((ListSerializer) descriptor.getSerializer()).getElementSerializer()));
	}

	@Override
	public  SubKeyedMapState getSubKeyedMapState(
		MapStateDescriptor descriptor
	) throws Exception {
		if (namespaceSerializer == null) {
			throw new RuntimeException("The namespace serializer has not been initialized.");
		}

		descriptor.initializeSerializerUnlessSet(operator.getExecutionConfig());
		MapSerializer mapSerializer = (MapSerializer) descriptor.getSerializer();
		return operator.getSubKeyedState(new SubKeyedMapStateDescriptor<>(
			descriptor.getName(),
			(TypeSerializer) operator.getKeySerializer(),
			(TypeSerializer) namespaceSerializer,
			mapSerializer.getKeySerializer(),
			mapSerializer.getValueSerializer()));
	}

	@Override
	public  SubKeyedSortedMapState getSubKeyedSortedMapState(
		SortedMapStateDescriptor descriptor
	) throws Exception {
		if (namespaceSerializer == null) {
			throw new RuntimeException("The namespace serializer has not been initialized.");
		}
		descriptor.initializeSerializerUnlessSet(operator.getExecutionConfig());
		SortedMapSerializer sortedMapSerializer = (SortedMapSerializer) descriptor.getSerializer();
		return operator.getSubKeyedState(new SubKeyedSortedMapStateDescriptor<>(
			descriptor.getName(),
			(TypeSerializer) operator.getKeySerializer(),
			(TypeSerializer) namespaceSerializer,
			sortedMapSerializer.getComparator(),
			sortedMapSerializer.getKeySerializer(),
			sortedMapSerializer.getValueSerializer()));
	}

	@Override
	public  StateMapView getStateMapView(
		String stateName,
		MapViewTypeInfo mapViewTypeInfo,
		boolean hasNamespace) throws Exception {

		MapStateDescriptor mapStateDescriptor = new MapStateDescriptor<>(
			stateName,
			mapViewTypeInfo.keyType(),
			mapViewTypeInfo.valueType());

		ValueStateDescriptor nullStateDescriptor = new ValueStateDescriptor<>(
			stateName + NULL_STATE_POSTFIX,
			mapViewTypeInfo.valueType());

		if (hasNamespace) {
			SubKeyedMapState mapState = getSubKeyedMapState(mapStateDescriptor);
			if (mapViewTypeInfo.nullAware()) {
				SubKeyedValueState nullState = getSubKeyedValueState(nullStateDescriptor);
				return new NullAwareSubKeyedStateMapView<>(mapState, nullState);
			} else {
				return new SubKeyedStateMapView<>(mapState);
			}
		} else {
			KeyedMapState mapState = getKeyedMapState(mapStateDescriptor);
			if (mapViewTypeInfo.nullAware()) {
				KeyedValueState nullState = getKeyedValueState(nullStateDescriptor);
				return new NullAwareKeyedStateMapView<>(mapState, nullState);
			} else {
				return new KeyedStateMapView<>(mapState);
			}
		}
	}

	@Override
	public  StateSortedMapView getStateSortedMapView(
		String stateName,
		SortedMapViewTypeInfo sortedMapViewTypeInfo,
		boolean hasNamespace) throws Exception {

		SortedMapStateDescriptor sortedMapStateDesc = new SortedMapStateDescriptor<>(
			stateName,
			sortedMapViewTypeInfo.comparator,
			sortedMapViewTypeInfo.keyType,
			sortedMapViewTypeInfo.valueType);

		if (!hasNamespace) {
			KeyedSortedMapState mapState = getKeyedSortedMapState(sortedMapStateDesc);
			return new KeyedStateSortedMapView<>(mapState);
		} else {
			throw new UnsupportedOperationException("SubKeyedState SortedMapView is not supported currently");
		}
	}

	@Override
	public  StateListView getStateListView(
		String stateName,
		ListViewTypeInfo listViewTypeInfo,
		boolean hasNamespace) throws Exception {

		ListStateDescriptor listStateDesc = new ListStateDescriptor<>(
			stateName,
			listViewTypeInfo.elementType());

		if (hasNamespace) {
			SubKeyedListState listState = getSubKeyedListState(listStateDesc);
			return new SubKeyedStateListView<>(listState);
		} else {
			KeyedListState listState = getKeyedListState(listStateDesc);
			return new KeyedStateListView<>(listState);
		}
	}

	@Override
	public void registerStateDataView(StateDataView stateDataView) {
		registeredStateDataViews.add(stateDataView);
	}

	@Override
	public  TypeSerializer getKeySerializer() {
		return (TypeSerializer) operator.getKeySerializer();
	}

	@Override
	public BaseRow currentKey() {
		return (BaseRow) operator.getCurrentKey();
	}

	@Override
	public void setCurrentKey(BaseRow key) {
		operator.setCurrentKey(key);
		// set current key to all the registered stateDataviews
		for (StateDataView dataView : registeredStateDataViews) {
			dataView.setCurrentKey(key);
		}
	}

	@Override
	public RuntimeContext getRuntimeContext() {
		return runtimeContext;
	}
}