org.apache.flink.runtime.state.gemini.engine.hashtable.GRegionKSortedMapImpl 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.engine.hashtable;

import org.apache.flink.runtime.state.gemini.engine.GRegionContext;
import org.apache.flink.runtime.state.gemini.engine.GeminiKSortedMap;
import org.apache.flink.runtime.state.gemini.engine.memstore.GSValue;
import org.apache.flink.runtime.state.gemini.engine.memstore.GSValueMap;
import org.apache.flink.runtime.state.gemini.engine.memstore.WriteBufferKSortedMapHashImpl;
import org.apache.flink.runtime.state.gemini.engine.page.GValueType;
import org.apache.flink.runtime.state.gemini.engine.page.PageIndex;
import org.apache.flink.runtime.state.gemini.engine.page.PageSerdeFlink2Key;
import org.apache.flink.runtime.state.gemini.engine.page.PageStoreHashKSortedMapImpl;
import org.apache.flink.runtime.state.gemini.engine.page.bmap.GComparator;
import org.apache.flink.runtime.util.EmptyIterator;
import org.apache.flink.util.Preconditions;

import org.apache.flink.shaded.guava18.com.google.common.collect.Maps;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;

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

/**
 * GRegionKSortedMapImpl.
 * K is type of key.
 * V is type of value.
 */

public class GRegionKSortedMapImpl extends AbstractGRegionKMapImpl> implements GeminiKSortedMap {

	private GComparator comparator;
	public GRegionKSortedMapImpl(GRegionContext gRegionContext) {
		super(gRegionContext, null);
	}

	public GRegionKSortedMapImpl(GRegionContext gRegionContext, PageIndex pageIndex) {
		super(gRegionContext, pageIndex);
	}

	@Override
	void init(PageIndex pageIndex) {
		PageSerdeFlink2Key pageSerdeFlink = (PageSerdeFlink2Key) gRegionContext.getPageSerdeFlink();
		Preconditions.checkArgument(pageSerdeFlink.getMapComparator() != null, "no MapComparator,error type");
		this.pageStore = new PageStoreHashKSortedMapImpl<>(this, pageIndex, regionEventExecutor);
		this.comparator = pageSerdeFlink.getMapComparator();
		this.writeBuffer = new WriteBufferKSortedMapHashImpl<>(this,
			regionEventExecutor,
			pageStore,
			comparator);
	}

	@Override
	public Map.Entry firstEntry(K key) {
		gContext.checkDBStatus();
		gContext.incAccessNumber();
		GSValueMap mapFromWriteBuffer = writeBuffer.get(key);
		Map> mapFromPageStore = pageStore.get(key);

		TreeMap map = mergeTwoMap(mapFromWriteBuffer, mapFromPageStore, false);

		return map == null ? null : copyEntry(map.firstEntry());
	}

	@Override
	public Map.Entry lastEntry(K key) {
		gContext.checkDBStatus();
		gContext.incAccessNumber();
		GSValueMap mapFromWriteBuffer = writeBuffer.get(key);
		Map> mapFromPageStore = pageStore.get(key);

		TreeMap map = mergeTwoMap(mapFromWriteBuffer, mapFromPageStore, false);
		return map == null ? null : copyEntry(map.lastEntry());
	}

	@Override
	public Iterator> headIterator(K key, MK endMapKey) {
		gContext.checkDBStatus();
		gContext.incAccessNumber();
		GSValueMap mapFromWriteBuffer = writeBuffer.get(key);
		Map> mapFromPageStore = pageStore.get(key);

		TreeMap map = mergeTwoMap(mapFromWriteBuffer, mapFromPageStore, true);

		//Empty Iterator, that's strange, let's ignore this now.
		return map == null ? EmptyIterator.get() : map.headMap(endMapKey).entrySet().iterator();
	}

	@Override
	public Iterator> tailIterator(K key, MK startMapKey) {
		gContext.checkDBStatus();
		gContext.incAccessNumber();
		GSValueMap mapFromWriteBuffer = writeBuffer.get(key);
		Map> mapFromPageStore = pageStore.get(key);

		TreeMap map = mergeTwoMap(mapFromWriteBuffer, mapFromPageStore, true);

		//Empty Iterator, that's strange, let's ignore this now.
		return map == null ? EmptyIterator.get() : map.tailMap(startMapKey).entrySet().iterator();
	}

	@Override
	public Iterator> subIterator(K key, MK startMapKey, MK endMapKey) {
		gContext.checkDBStatus();
		gContext.incAccessNumber();
		GSValueMap mapFromWriteBuffer = writeBuffer.get(key);
		Map> mapFromPageStore = pageStore.get(key);

		TreeMap map = mergeTwoMap(mapFromWriteBuffer, mapFromPageStore, true);

		//Empty Iterator, that's strange, let's ignore this now.
		return map == null ? EmptyIterator.get() : map.subMap(startMapKey, endMapKey).entrySet().iterator();
	}

	@Override
	public void put(K key, SortedMap value) {
		gContext.checkDBStatus();
		gContext.incAccessNumber();
		long seqID = gRegionContext.getNextSeqID();
		writeBuffer.put(key, Maps.transformEntries(value, (k, v) -> GSValue.of(v, GValueType.PutValue, seqID)));
	}

	@Override
	public SortedMap get(K key) {
		gContext.checkDBStatus();
		gContext.incAccessNumber();
		return internalGet(key, true);
	}

	@Override
	public SortedMap getOrDefault(K key, SortedMap defaultValue) {
		gContext.checkDBStatus();
		SortedMap ret = get(key);
		return (ret == null || ret.isEmpty()) ? defaultValue : ret;
	}

	@Override
	public Map> getAll() {
		gContext.checkDBStatus();
		gContext.incAccessNumber();
		Map> results = new HashMap<>();
		getAll(results);
		return results;
	}

	private SortedMap genMapFromGSValueMap(GSValueMap gsValueMap) {
		return genMapFromGSValueMap(gsValueMap, false);
	}

	private SortedMap genMapFromGSValueMap(GSValueMap gsValueMap, boolean checkReadCopy) {
		SortedMap> valueMap = (SortedMap>) gsValueMap.getValue();
		SortedMap ret = new TreeMap<>(comparator.getJDKCompactor());
		if (valueMap != null) {
			for (SortedMap.Entry> entry : valueMap.entrySet()) {
				GSValue gsValue = entry.getValue();
				if (GValueType.Delete.equals(gsValue.getValueType()) || gRegionContext.filterState(gsValue.getSeqID())) {
					continue;
				}
				MK mk = entry.getKey();
				MV mv = entry.getValue().getValue();
				ret.put(checkReadCopy ? copyMKIfNeeded(mk) : mk, checkReadCopy ? copyMVIfNeeded(mv) : mv);
			}
		}
		return ret.isEmpty() ? null : ret;
	}

	private TreeMap mergeTwoMap(
		GSValueMap mapFromWriteBuffer,
		Map> mapFromPageStore) {
		return mergeTwoMap(mapFromWriteBuffer, mapFromPageStore, false);
	}
	private TreeMap mergeTwoMap(
		GSValueMap mapFromWriteBuffer,
		Map> mapFromPageStore,
		boolean checkReadCopy) {
		TreeMap map = new TreeMap<>(comparator.getJDKCompactor());
		if (mapFromPageStore != null) {
			for (Map.Entry> entry : mapFromPageStore.entrySet()) {
				if (gRegionContext.filterState(entry.getValue().getSeqID())) {
					continue;
				}
				if (entry.getValue().getValueType() == GValueType.PutValue) {
					map.put(entry.getKey(), entry.getValue().getValue());
				} else {
					checkState(entry.getValue().getValueType().equals(GValueType.Delete));
				}
			}
		}

		if (mapFromWriteBuffer != null) {
			for (Map.Entry> entry : mapFromWriteBuffer.getValue().entrySet()) {
				if (gRegionContext.filterState(entry.getValue().getSeqID())) {
					continue;
				}
				if (entry.getValue().getValueType().equals(GValueType.PutValue)) {
					MK mk = entry.getKey();
					MV mv = entry.getValue().getValue();
					map.put(checkReadCopy ? copyMKIfNeeded(mk) : mk, checkReadCopy ? copyMVIfNeeded(mv) : mv);
				} else {
					checkState(entry.getValue().getValueType().equals(GValueType.Delete));
					map.remove(entry.getKey());
				}
			}
		}
		return map == null || map.isEmpty() ? null : map;
	}

	@Override
	public void getAll(Map> results) {
		gContext.checkDBStatus();
		//TODO provide better performance.
		gContext.incAccessNumber();
		Set allKeysIncludeDelete = new HashSet<>();
		writeBuffer.allKeysIncludeDeleted(allKeysIncludeDelete);
		pageStore.allKeysIncludeDeleted(allKeysIncludeDelete);

		for (K key : allKeysIncludeDelete) {
			SortedMap mapResult = get(key);

			if (mapResult != null) {
				results.put(copyKeyIfNeeded(key), mapResult);
			}
		}
	}

	@Override
	protected SortedMap internalGet(K key, boolean checkReadCopy) {
		GSValueMap map = writeBuffer.get(key);
		if (map != null) {
			switch (map.getValueType()) {
				case Delete:
					return null;
				case PutMap:
					return genMapFromGSValueMap(map, checkReadCopy);
				default:
					// nothing todo
			}
		}

		Map> mapFromPageStore = pageStore.get(key);
		SortedMap ret = mergeTwoMap(map, mapFromPageStore, checkReadCopy);
		return ret == null || ret.isEmpty() ? null : ret;
	}

	private Map.Entry copyEntry(Map.Entry entry) {
		if (entry == null || !readCopy) {
			return entry;
		}
		return new Map.Entry() {

			private MK mk = mkSerializer.copy(entry.getKey());

			private MV mv = mvSerializer.copy(entry.getValue());

			@Override
			public MK getKey() {
				return mk;
			}

			@Override
			public MV getValue() {
				return mv;
			}

			@Override
			public MV setValue(MV value) {
				MV oldValue = mv;
				mv = value;
				return oldValue;
			}
		};
	}
}