org.apache.flink.runtime.state.gemini.engine.page.bmap.GBinarySortedMap 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.page.bmap;

import org.apache.flink.annotation.VisibleForTesting;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.runtime.state.gemini.engine.GRegionContext;
import org.apache.flink.runtime.state.gemini.engine.exceptions.GeminiRuntimeException;
import org.apache.flink.runtime.state.gemini.engine.filter.StateFilter;
import org.apache.flink.runtime.state.gemini.engine.memstore.GSValue;
import org.apache.flink.runtime.state.gemini.engine.page.DataPage;
import org.apache.flink.runtime.state.gemini.engine.page.GValueType;
import org.apache.flink.runtime.state.gemini.engine.page.compress.GCompressAlgorithm;
import org.apache.flink.runtime.state.gemini.engine.rm.Allocator;
import org.apache.flink.runtime.state.gemini.engine.rm.GByteBuffer;
import org.apache.flink.runtime.state.gemini.engine.rm.GUnPooledByteBuffer;
import org.apache.flink.util.MathUtils;

import javax.annotation.Nullable;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicReference;

/**
 * GBinarySortedMap for map type of Value.
 * TODO LIMITATION: now if do remove or add opertor to the map instance returned by head/tail/sub operator,
 * it don't affect the origin map.
 */
public class GBinarySortedMap extends GAbstractSortedMap {
	public static final GBinarySortedMap EMPTY_G_BINARY_SORTEDMAP = new GBinarySortedMap(null, null, null, null);

	private final GBinaryHashMap gBinaryHashMap;
	private final ByteBuffer data;
	private final GSortedHeaderImpl gSortedHeader;
	private final GComparator gComparator;
	private final TypeSerializer keyTypeSerializer;

	public GBinarySortedMap(
		ByteBuffer data, TypeSerializer keyTypeSerializer, GComparator gComparator) {
		int indexLen = data == null ? 0 : GHashHeaderImpl.getHeaderIndexCount(data);
		GSortedHeaderImpl gSortedHeaderImpl = GSortedHeaderImpl.getPageHelper(indexLen);
		this.gSortedHeader = gSortedHeaderImpl;
		this.data = (data != null && data.capacity() == 0) ? null : data;
		this.keyTypeSerializer = keyTypeSerializer;
		this.gComparator = gComparator;
		this.gBinaryHashMap = new GBinaryHashMap<>(gSortedHeaderImpl, new GUnPooledByteBuffer(data), keyTypeSerializer);
	}

	public GBinarySortedMap(
		GSortedHeaderImpl gSortedHeader,
		ByteBuffer data,
		TypeSerializer keyTypeSerializer,
		GComparator gComparator) {
		this.data = (data != null && data.capacity() == 0) ? null : data;
		this.gSortedHeader = gSortedHeader;
		this.keyTypeSerializer = keyTypeSerializer;
		this.gComparator = gComparator;
		this.gBinaryHashMap = new GBinaryHashMap<>(gSortedHeader, new GUnPooledByteBuffer(data), keyTypeSerializer);
	}

	@Override
	public BinaryValue get(Object key) {
		return this.gBinaryHashMap.get(key);
	}

	@Override
	public int size() {
		return this.gBinaryHashMap.size();
	}

	@Override
	public int bytesSize() {
		return this.gBinaryHashMap.bytesSize();
	}

	@Override
	public boolean isEmpty() {
		return this.gBinaryHashMap.isEmpty();
	}

	@Override
	public boolean containsKey(Object key) {
		return get(key) != null;
	}

	@Override
	public long getCompactionCount() {
		return this.gBinaryHashMap.getCompactionCount();
	}

	@Override
	public int keyCount() {
		return this.gBinaryHashMap.keyCount();
	}

	public byte[] getDataByte(boolean checksumEnable) {
		return this.gBinaryHashMap.getDataByte(checksumEnable);
	}

	public int getLogicPageId() {
		return this.gBinaryHashMap.getLogicPageId();
	}

	@Override
	public Map getBinaryMap() {
		return this.gBinaryHashMap.getBinaryMap();
	}

	@Override
	public SortedMap getSortedBinaryMap() {
		int keyCount = keyCount();
		TreeMap dataMap = new TreeMap<>(gComparator.getJDKBinaryCompactor());
		if (keyCount == 0) {
			return dataMap;
		}
		int indexLen = this.gBinaryHashMap.indexCount();
		int sortedIndexBaseOffset = gSortedHeader.getSortedIndexBaseOffset(indexLen, keyCount);
		int baseKeyOffset = GHashHeaderImpl.getHeaderKeyOffset(data);
		int baseValueOffset = GHashHeaderImpl.getHeaderValueOffset(data);
		AtomicReference uncompressKeyData = new AtomicReference<>();
		AtomicReference uncompressValueData = new AtomicReference<>();
		for (int k = 0; k < keyCount; k++) {
			//get key slot.
			int slotNum = gSortedHeader.getSortedIndexBySlot(data, sortedIndexBaseOffset, k);
			//BinaryKey
			BinaryKey binaryKey = this.gBinaryHashMap.getBinaryKey(keyCount,
				indexLen,
				baseKeyOffset,
				baseValueOffset,
				slotNum,
				uncompressKeyData);
			//BinaryValue
			BinaryValue binaryValue = this.gBinaryHashMap.getBinaryValue(keyCount,
				indexLen,
				baseValueOffset,
				slotNum,
				uncompressValueData);
			dataMap.put(binaryKey, binaryValue);
		}

		return dataMap;
	}

	@Override
	public TypeSerializer getKeyTypeSerializer() {
		return this.keyTypeSerializer;
	}

	@Override
	public Comparator comparator() {
		return this.gComparator.getJDKCompactor();
	}

	@Override
	public SortedMap subMap(MK fromKey, MK toKey) {
		int keyCount = keyCount();
		TreeMap dataMap = new TreeMap<>(gComparator.getJDKCompactor());
		if (keyCount == 0) {
			return dataMap;
		}

		if (fromKey != null && toKey != null && gComparator.compare(fromKey, toKey) > 0) {
			throw new IllegalArgumentException("fromKey > toKey");
		}

		int startIndexSlot = findSortedIndexSlotHigher(fromKey, keyCount, true);
		int endIndexSlot = findSortedIndexSlotLower(toKey, keyCount, false);
		if (endIndexSlot < startIndexSlot || startIndexSlot >= keyCount) {
			return dataMap;
		} else if (endIndexSlot >= keyCount) {
			throw new GeminiRuntimeException("Internal Bug");
		}
		int indexLen = this.gBinaryHashMap.indexCount();
		int sortedIndexBaseOffset = gSortedHeader.getSortedIndexBaseOffset(indexLen, keyCount);
		int baseKeyOffset = GHashHeaderImpl.getHeaderKeyOffset(data);
		int baseValueOffset = GHashHeaderImpl.getHeaderValueOffset(data);
		AtomicReference uncompressData = new AtomicReference<>();
		AtomicReference uncompressValueData = new AtomicReference<>();
		for (int k = startIndexSlot; k <= endIndexSlot; k++) {
			int slotNum = gSortedHeader.getSortedIndexBySlot(data, sortedIndexBaseOffset, k);
			MK key = this.gBinaryHashMap.getKeyBySlot(keyCount,
				baseKeyOffset,
				baseValueOffset,
				slotNum,
				uncompressData);
			//BinaryValue
			BinaryValue binaryValue = this.gBinaryHashMap.getBinaryValue(keyCount,
				indexLen,
				baseValueOffset,
				slotNum,
				uncompressValueData);
			dataMap.put(key, binaryValue);
		}
		return dataMap;
	}

	private int findSortedIndexSlotLower(MK toKey, int keyCount, boolean inclusive) {
		int low = 0;
		int high = keyCount - 1;

		if (toKey == null) {
			return high;
		}

		int indexLen = this.gBinaryHashMap.indexCount();
		int sortedIndexBaseOffset = gSortedHeader.getSortedIndexBaseOffset(indexLen, keyCount);
		int baseKeyOffset = GHashHeaderImpl.getHeaderKeyOffset(data);
		int baseValueOffset = GHashHeaderImpl.getHeaderValueOffset(data);
		AtomicReference uncompressData = new AtomicReference<>();

		while (low <= high) {
			int mid = (low + high) >>> 1;
			int keySlotNum = gSortedHeader.getSortedIndexBySlot(data, sortedIndexBaseOffset, mid);
			MK midVal = this.gBinaryHashMap.getKeyBySlot(keyCount,
				baseKeyOffset,
				baseValueOffset,
				keySlotNum,
				uncompressData);

			if (gComparator.compare(midVal, toKey) < 0) {
				low = mid + 1;
			} else if (gComparator.compare(midVal, toKey) > 0) {
				high = mid - 1;
			} else {
				return inclusive ? mid : mid - 1; // key found
			}
		}
		return low - 1;  // key not found.
	}

	private int findSortedIndexSlotHigher(MK fromKey, int keyCount, boolean inclusive) {
		if (fromKey == null) {
			return 0;
		}
		int low = 0;
		int high = keyCount - 1;

		int indexLen = this.gBinaryHashMap.indexCount();
		int sortedIndexBaseOffset = gSortedHeader.getSortedIndexBaseOffset(indexLen, keyCount);
		int baseKeyOffset = GHashHeaderImpl.getHeaderKeyOffset(data);
		int baseValueOffset = GHashHeaderImpl.getHeaderValueOffset(data);
		AtomicReference uncompressData = new AtomicReference<>();

		while (low <= high) {
			int mid = (low + high) >>> 1;
			int keySlotNum = gSortedHeader.getSortedIndexBySlot(data, sortedIndexBaseOffset, mid);
			MK midVal = this.gBinaryHashMap.getKeyBySlot(keyCount,
				baseKeyOffset,
				baseValueOffset,
				keySlotNum,
				uncompressData);

			if (gComparator.compare(midVal, fromKey) < 0) {
				low = mid + 1;
			} else if (gComparator.compare(midVal, fromKey) > 0) {
				high = mid - 1;
			} else {
				return inclusive ? mid : mid + 1; // key found
			}
		}
		return low;  // key not found.
	}

	@Override
	public SortedMap headMap(MK toKey) {
		return subMap(null, toKey);
	}

	@Override
	public SortedMap tailMap(MK fromKey) {
		return subMap(fromKey, null);
	}

	@Override
	public MK firstKey() {
		int keyCount = keyCount();
		if (keyCount == 0) {
			return null;
		}
		return getMkBySortedIndex(keyCount, 0);
	}

	@Override
	public MK lastKey() {
		int keyCount = keyCount();
		if (keyCount == 0) {
			return null;
		}
		return getMkBySortedIndex(keyCount, keyCount - 1);
	}

	private MK getMkBySortedIndex(int keyCount, int sortedIndexSlot) {

		int indexLen = this.gBinaryHashMap.indexCount();
		int sortedIndexBaseOffset = gSortedHeader.getSortedIndexBaseOffset(indexLen, keyCount);
		int baseKeyOffset = GHashHeaderImpl.getHeaderKeyOffset(data);
		int baseValueOffset = GHashHeaderImpl.getHeaderValueOffset(data);
		int keySlotNum = gSortedHeader.getSortedIndexBySlot(data, sortedIndexBaseOffset, sortedIndexSlot);
		return this.gBinaryHashMap.getKeyBySlot(keyCount, baseKeyOffset, baseValueOffset, keySlotNum, null);
	}

	public long getVersion() {
		return this.gBinaryHashMap.getVersion();
	}

	public ByteBuffer getData() {
		return this.gBinaryHashMap.getData();
	}

	public static  GBinarySortedMap of(
		DataPage.DataPageType dataPageType,
		List>> keyValueList,
		TypeSerializer keySerializer,
		TypeSerializer valueSerializer,
		GComparator gComparator,
		long version,
		int logicPageId,
		Allocator allocator,
		long compactionCount) {

		int totalKeys = keyValueList.size();
		if (totalKeys == 0) {
			return EMPTY_G_BINARY_SORTEDMAP;
		}

		int indexLen = MathUtils.roundUpToPowerOfTwo(totalKeys);
		GSortedHeaderImpl gheader = GSortedHeaderImpl.getPageHelper(indexLen);

		GByteBuffer gByteBuffer = GBinaryHashMap.genByteBuffer(dataPageType,
			keyValueList,
			indexLen,
			totalKeys,
			gheader,
			keySerializer,
			valueSerializer,
			version,
			logicPageId,
			allocator,
			compactionCount,
			//sortedMap is a sigle value. no need to compress.
			GCompressAlgorithm.None);

		//do write sorted index
		//1. dataSet have ordered by hashCode, and it's order is same with what they are been stored.
		List> slotAndKeyList = new ArrayList<>();
		int i = 0;
		for (Tuple2> item : keyValueList) {
			slotAndKeyList.add(Tuple2.of(item.f0, i));
			i++;
		}
		Collections.sort(slotAndKeyList, (o1, o2) -> gComparator.compare(o1.f0, o2.f0));
		int index = 0;
		int baseOffset = gheader.getSortedIndexBaseOffset(indexLen, totalKeys);
		for (Tuple2 slot : slotAndKeyList) {
			gheader.writeSortedIndexBySlot(gByteBuffer.getByteBuffer(), baseOffset, index, slot.f1);
			index++;
		}

		return new GBinarySortedMap<>(gheader, gByteBuffer.getByteBuffer(), keySerializer, gComparator);
	}

	public static  GBinarySortedMap ofBinaryList(
		DataPage.DataPageType dataPageType,
		boolean isMajor,
		long version,
		int logicPageId,
		TypeSerializer keyTypeSerializer,
		GComparator gComparator,
		Allocator allocator,
		Map newMap,
		long compactionCount,
		@Nullable StateFilter stateFilter,
		@Nullable GRegionContext gRegionContext) {
		List> dataSet = new ArrayList<>();

		for (Map.Entry entry : newMap.entrySet()) {
			BinaryValue binaryValue = entry.getValue();
			if (isMajor && (binaryValue.getgValueType() == GValueType.Delete || (stateFilter != null && stateFilter.filter(
				gRegionContext,
				binaryValue.getSeqID())))) {
				continue;
			}
			dataSet.add(Tuple2.of(entry.getKey(), entry.getValue()));
		}

		int totalKeys = dataSet.size();
		if (totalKeys == 0) {
			return EMPTY_G_BINARY_SORTEDMAP;
		}
		int indexLen = MathUtils.roundUpToPowerOfTwo(totalKeys);
		GSortedHeaderImpl gheader = GSortedHeaderImpl.getPageHelper(indexLen);
		GByteBuffer byteBuffer = GBinaryHashMap.genByteBufferByBinary(dataPageType,
			dataSet,
			indexLen,
			totalKeys,
			gheader,
			version,
			logicPageId,
			allocator,
			compactionCount,
			//sortedMap is a sigle value. no need to compress.
			GCompressAlgorithm.None);

		//do write sorted index
		//1. dataSet have ordered by hashCode, and it's order is same with what they are been stored.
		List> slotAndKeyList = new ArrayList<>();
		int i = 0;
		for (Tuple2 item : dataSet) {
			slotAndKeyList.add(Tuple2.of(item.f0, i));
			i++;
		}
		Collections.sort(slotAndKeyList, (o1, o2) -> gComparator.compare(o1.f0, o2.f0));
		int index = 0;
		int baseOffset = gheader.getSortedIndexBaseOffset(indexLen, totalKeys);
		for (Tuple2 slot : slotAndKeyList) {
			gheader.writeSortedIndexBySlot(byteBuffer.getByteBuffer(), baseOffset, index, slot.f1);
			index++;
		}
		return new GBinarySortedMap<>(gheader, byteBuffer.getByteBuffer(), keyTypeSerializer, gComparator);
	}

	@VisibleForTesting
	public GBinaryHashMap getgBinaryHashMap() {
		return gBinaryHashMap;
	}

}