org.apache.flink.runtime.state.gemini.subkeyed.GeminiSubKeyedListStateImpl 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.flink.runtime.state.gemini.subkeyed;

import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.runtime.state.StateStorage;
import org.apache.flink.runtime.state.gemini.engine.GRegion;
import org.apache.flink.runtime.state.gemini.engine.GeminiPKey2;
import org.apache.flink.runtime.state.gemini.engine.hashtable.GRegionKListImpl;
import org.apache.flink.runtime.state.gemini.engine.hashtable.GRegionKMapImpl;
import org.apache.flink.runtime.state.gemini.engine.hashtable.GTableSubKeyedListImpl;
import org.apache.flink.runtime.state.subkeyed.SubKeyedListState;
import org.apache.flink.runtime.state.subkeyed.SubKeyedListStateDescriptor;

import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;
import java.util.stream.Collectors;

import static org.apache.flink.util.Preconditions.checkNotNull;

/**
 * An implementation of {@link SubKeyedListState} backed by a state storage.
 * The pairs in the storage are formatted as {(K, N) -> List{E}}.
 * Because the pairs are partitioned by K, all the elements under the same key
 * reside in the same group. They can be easily retrieved with a prefix iterator
 * on the key.
 *
 * @param  Type of the keys in the state.
 * @param  Type of the namespaces in the state.
 * @param  Type of the elements in the state.
 */
public final class GeminiSubKeyedListStateImpl implements SubKeyedListState {

	/**
	 * The descriptor of this state.
	 */
	private final SubKeyedListStateDescriptor descriptor;
	private final GTableSubKeyedListImpl table;
	private final Byte dummyByte = new Byte((byte) 0);

	//--------------------------------------------------------------------------

	/**
	 * Constructor with the state storage to store elements.
	 *
	 * @param descriptor The state storage where elements are stored.
	 * @param table
	 */
	public GeminiSubKeyedListStateImpl(
		SubKeyedListStateDescriptor descriptor,
		GTableSubKeyedListImpl table) {
		this.table = checkNotNull(table);
		this.descriptor = checkNotNull(descriptor);
	}

	@Override
	public SubKeyedListStateDescriptor getDescriptor() {
		return descriptor;
	}

	//--------------------------------------------------------------------------

	@Override
	public boolean contains(K key, N namespace) {
		if (key == null || namespace == null) {
			return false;
		}

		return getRegion(key, namespace).contains(table.getGeminiPKey2(key, namespace));
	}

	@Override
	public List get(K key, N namespace) {
		return getOrDefault(key, namespace, null);
	}

	@Override
	public List getOrDefault(K key, N namespace, List defaultList) {
		if (key == null || namespace == null) {
			return defaultList;
		}

		return getRegion(key, namespace).getOrDefault(table.getGeminiPKey2(key, namespace), defaultList);
	}

	@SuppressWarnings("unchecked")
	@Override
	public Map> getAll(K key) {
		if (key == null) {
			return Collections.emptyMap();
		}

		Map> result = new HashMap<>();
		Set visitedNamespace = new HashSet<>();
		Iterator namepsaceIter = table.getSecondaryKeyByFirstKey(key);
		while (namepsaceIter.hasNext()) {
			N namespace = namepsaceIter.next();
			if (visitedNamespace.contains(namespace)) {
				continue;
			}
			List value = get(key, namespace);
			visitedNamespace.add(namespace);
			result.put(namespace, value);
		}
		return result;
	}

	@Override
	public void add(K key, N namespace, E element) {
		checkNotNull(key);
		checkNotNull(namespace);
		checkNotNull(element);

		getRegion(key, namespace).add(table.getGeminiPKey2(key, namespace), element);
		getIndexRegion(key, namespace).add(key, namespace, dummyByte);
	}

	@Override
	public void addAll(K key, N namespace, Collection elements) {
		checkNotNull(key);
		checkNotNull(namespace);
		checkNotNull(elements);

		if (elements.isEmpty()) {
			return;
		}

		getRegion(key, namespace).addAll(table.getGeminiPKey2(key, namespace), elements);
		getIndexRegion(key, namespace).add(key, namespace, dummyByte);
	}

	@Override
	public void put(K key, N namespace, E element) {
		checkNotNull(key);
		checkNotNull(namespace);
		checkNotNull(element);

		// TODO: add a function put(K, N, E) in GRegionKListImpl
		getRegion(key, namespace).put(table.getGeminiPKey2(key, namespace), Arrays.asList(element));
		getIndexRegion(key, namespace).add(key, namespace, dummyByte);
	}

	@Override
	public void putAll(K key, N namespace, Collection elements) {
		checkNotNull(key);
		checkNotNull(namespace);
		checkNotNull(elements);

		if (elements.isEmpty()) {
			return;
		}

		getRegion(key, namespace).put(table.getGeminiPKey2(key, namespace), elements.stream().collect(Collectors.toList()));
		getIndexRegion(key, namespace).add(key, namespace, dummyByte);
	}

	@Override
	public void remove(K key, N namespace) {
		if (key == null || namespace == null) {
			return;
		}

		getRegion(key, namespace).remove(table.getGeminiPKey2(key, namespace));
		getIndexRegion(key, namespace).remove(key, namespace);
	}

	@Override
	public boolean remove(K key, N namespace, E elementToRemove) {
		if (key == null || namespace == null) {
			return false;
		}

		getRegion(key, namespace).remove(table.getGeminiPKey2(key, namespace), elementToRemove);
		// do not update index now.
		return true;
	}

	@Override
	public boolean removeAll(K key, N namespace, Collection elements) {
		if (key == null || namespace == null || elements == null || elements.isEmpty()) {
			return false;
		}

		getRegion(key, namespace).removeAll(table.getGeminiPKey2(key, namespace), elements);
		// do not update index now.
		return true;
	}

	@Override
	public void removeAll(K key) {
		if (key == null) {
			return;
		}

		Iterator namespaceIter = table.getSecondaryKeyByFirstKey(key);
		while (namespaceIter.hasNext()) {
			N namespace = namespaceIter.next();
			remove(key, namespace);
		}
	}

	@Override
	public Iterator iterator(K key) {
		checkNotNull(key);

		Iterator namespaceIter = table.getSecondaryKeyByFirstKey(key);
		return new Iterator() {
			N nextNamespace = null;
			N lastNamespace = null;
			@Override
			public boolean hasNext() {
				while (nextNamespace == null && namespaceIter.hasNext()) {
					nextNamespace = namespaceIter.next();
					if (!contains(key, nextNamespace)) {
						// set the flag to move on.
						nextNamespace = null;
					}
				}
				return nextNamespace != null;
			}

			@Override
			public N next() {
				if (!hasNext()) {
					throw new NoSuchElementException();
				}

				lastNamespace = nextNamespace;
				nextNamespace = null;
				return lastNamespace;
			}

			@Override
			public void remove() {
				GeminiSubKeyedListStateImpl.this.remove(key, lastNamespace);
			}
		};
	}

	@Override
	public StateStorage> getStateStorage() {
		throw new UnsupportedOperationException();
	}

	@Override
	public E poll(K key, N namespace) {
		return getRegion(key, namespace).poll(table.getGeminiPKey2(key, namespace));
	}

	@Override
	public E peek(K key, N namespace) {
		return getRegion(key, namespace).peek(table.getGeminiPKey2(key, namespace));
	}

	@Override
	public Iterable keys(N namespace) {
		Iterator iterator = table.regionIterator();
		Set result = new HashSet<>();
		while (iterator.hasNext()) {
			GRegionKListImpl, E> cur = (GRegionKListImpl, E>) iterator.next();
			result.addAll(
				cur.getAll().entrySet().stream()
					.filter(m -> namespace.equals(m.getKey().getSecondKey()))
					.map(y -> y.getKey().getFirstKey())
					.collect(Collectors.toSet()));
		}
		return result.isEmpty() ? null : result;
	}

	@Override
	public byte[] getSerializedValue(
		final byte[] serializedKeyAndNamespace,
		final TypeSerializer safeKeySerializer,
		final TypeSerializer safeNamespaceSerializer,
		final TypeSerializer> safeValueSerializer) {

		// TODO
		throw new UnsupportedOperationException();
	}

	@SuppressWarnings("unchecked")
	private GRegionKListImpl, E> getRegion(K key, N namespace) {
		return table.getRegion(table.getGeminiPKey2(key, namespace));
	}

	private GRegionKMapImpl getIndexRegion(K key, N namespace) {
		return table.getIndexRegion(table.getGeminiPKey2(key, namespace));
	}
}