tikv
/
client-java
mirror of https://github.com/tikv/client-java.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

apply comments 

Signed-off-by: birdstorm <samuelwyf@hotmail.com>
This commit is contained in:
birdstorm 2018-12-19 21:19:43 +08:00
parent 6aea6bbf2b
commit f964343f13
2 changed files with 21 additions and 26 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
src/main/java/org/tikv/raw/RawKVClient.java
View File

@ -89,20 +89,21 @@ public class RawKVClient implements AutoCloseable {
} }
/** /**
* Put a list of raw key-value pair to TiKV * Put a set of raw key-value pair to TiKV
* *
* @param keys keys * @param kvPairs kvPairs
* @param values values
*/ */
public void batchPut(List<ByteString> keys, List<ByteString> values) { public void batchPut(Map<ByteString, ByteString> kvPairs) {
batchPut(ConcreteBackOffer.newRawKVBackOff(), keys, values); batchPut(ConcreteBackOffer.newRawKVBackOff(), kvPairs);
} }
private void batchPut(BackOffer backOffer, List<ByteString> keys, List<ByteString> values) { private void batchPut(BackOffer backOffer, List<ByteString> keys, List<ByteString> values) {
assert keys.size() == values.size();
Map<ByteString, ByteString> keysToValues = mapKeysToValues(keys, values); Map<ByteString, ByteString> keysToValues = mapKeysToValues(keys, values);
Map<Long, List<ByteString>> groupKeys = groupKeysByRegion(keys); batchPut(backOffer, keysToValues);
keys.clear(); }
private void batchPut(BackOffer backOffer, Map<ByteString, ByteString> kvPairs) {
Map<Long, List<ByteString>> groupKeys = groupKeysByRegion(kvPairs.keySet());
List<Batch> batches = new ArrayList<>(); List<Batch> batches = new ArrayList<>();
for (Map.Entry<Long, List<ByteString>> entry : groupKeys.entrySet()) { for (Map.Entry<Long, List<ByteString>> entry : groupKeys.entrySet()) {
@ -110,7 +111,7 @@ public class RawKVClient implements AutoCloseable {
batches, batches,
entry.getKey(), entry.getKey(),
entry.getValue(), entry.getValue(),
entry.getValue().stream().map(keysToValues::get).collect(Collectors.toList()), entry.getValue().stream().map(kvPairs::get).collect(Collectors.toList()),
RAW_BATCH_PUT_SIZE); RAW_BATCH_PUT_SIZE);
} }
sendBatchPut(backOffer, batches); sendBatchPut(backOffer, batches);
@ -238,7 +239,7 @@ public class RawKVClient implements AutoCloseable {
* @param keys keys * @param keys keys
* @return a mapping of keys and their regionId * @return a mapping of keys and their regionId
*/ */
private Map<Long, List<ByteString>> groupKeysByRegion(List<ByteString> keys) { private Map<Long, List<ByteString>> groupKeysByRegion(Set<ByteString> keys) {
Map<Long, List<ByteString>> groups = new HashMap<>(); Map<Long, List<ByteString>> groups = new HashMap<>();
TiRegion lastRegion = null; TiRegion lastRegion = null;
for (ByteString key : keys) { for (ByteString key : keys) {

26
src/test/java/org/tikv/raw/RawKVClientTest.java
View File

@ -1,6 +1,5 @@
package org.tikv.raw; package org.tikv.raw;
import static org.tikv.raw.RawKVClient.mapKeysToValues;
import com.google.protobuf.ByteString; import com.google.protobuf.ByteString;
import java.util.*; import java.util.*;
@ -259,15 +258,13 @@ public class RawKVClientTest {
int i = cnt; int i = cnt;
completionService.submit( completionService.submit(
() -> { () -> {
List<ByteString> keyList = new ArrayList<>(); Map<ByteString, ByteString> kvPairs = new HashMap<>();
List<ByteString> valueList = new ArrayList<>();
for (int j = 0; j < base; j++) { for (int j = 0; j < base; j++) {
int num = i * base + j; int num = i * base + j;
ByteString key = orderedKeys.get(num), value = values.get(num); ByteString key = orderedKeys.get(num), value = values.get(num);
keyList.add(key); kvPairs.put(key, value);
valueList.add(value);
} }
client.batchPut(keyList, valueList); client.batchPut(kvPairs);
return null; return null;
}); });
} }
@ -282,15 +279,13 @@ public class RawKVClientTest {
+ " put=" + " put="
+ rawKeys().size()); + rawKeys().size());
} else { } else {
List<ByteString> keyList = new ArrayList<>(); Map<ByteString, ByteString> kvPairs = new HashMap<>();
List<ByteString> valueList = new ArrayList<>();
for (int i = 0; i < putCases; i++) { for (int i = 0; i < putCases; i++) {
ByteString key = randomKeys.get(i), value = values.get(r.nextInt(KEY_POOL_SIZE)); ByteString key = randomKeys.get(i), value = values.get(r.nextInt(KEY_POOL_SIZE));
data.put(key, value); data.put(key, value);
keyList.add(key); kvPairs.put(key, value);
valueList.add(value);
} }
checkBatchPut(keyList, valueList); checkBatchPut(kvPairs);
} }
} }
@ -402,11 +397,10 @@ public class RawKVClientTest {
assert client.get(key).equals(value); assert client.get(key).equals(value);
} }
private void checkBatchPut(List<ByteString> keyList, List<ByteString> valueList) { private void checkBatchPut(Map<ByteString, ByteString> kvPairs) {
client.batchPut(keyList, valueList); client.batchPut(kvPairs);
Map<ByteString, ByteString> keysToValues = mapKeysToValues(keyList, valueList); for (Map.Entry<ByteString, ByteString> kvPair : kvPairs.entrySet()) {
for (ByteString key : keyList) { assert client.get(kvPair.getKey()).equals(kvPair.getValue());
assert client.get(key).equals(keysToValues.get(key));
} }
} }