From f964343f132883058e0ba9d22632701a9cdb0d99 Mon Sep 17 00:00:00 2001
From: birdstorm <samuelwyf@hotmail.com>
Date: Wed, 19 Dec 2018 21:19:43 +0800
Subject: [PATCH] apply comments

Signed-off-by: birdstorm <samuelwyf@hotmail.com>
---
 src/main/java/org/tikv/raw/RawKVClient.java   | 21 ++++++++-------
 .../java/org/tikv/raw/RawKVClientTest.java    | 26 +++++++------------
 2 files changed, 21 insertions(+), 26 deletions(-)

diff --git a/src/main/java/org/tikv/raw/RawKVClient.java b/src/main/java/org/tikv/raw/RawKVClient.java
index 6f7159b6fb..056cf54f47 100644
--- a/src/main/java/org/tikv/raw/RawKVClient.java
+++ b/src/main/java/org/tikv/raw/RawKVClient.java
@@ -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 values values
+   * @param kvPairs kvPairs
    */
-  public void batchPut(List<ByteString> keys, List<ByteString> values) {
-    batchPut(ConcreteBackOffer.newRawKVBackOff(), keys, values);
+  public void batchPut(Map<ByteString, ByteString> kvPairs) {
+    batchPut(ConcreteBackOffer.newRawKVBackOff(), kvPairs);
   }
 
   private void batchPut(BackOffer backOffer, List<ByteString> keys, List<ByteString> values) {
-    assert keys.size() == values.size();
     Map<ByteString, ByteString> keysToValues = mapKeysToValues(keys, values);
-    Map<Long, List<ByteString>> groupKeys = groupKeysByRegion(keys);
-    keys.clear();
+    batchPut(backOffer, keysToValues);
+  }
+
+  private void batchPut(BackOffer backOffer, Map<ByteString, ByteString> kvPairs) {
+    Map<Long, List<ByteString>> groupKeys = groupKeysByRegion(kvPairs.keySet());
     List<Batch> batches = new ArrayList<>();
 
     for (Map.Entry<Long, List<ByteString>> entry : groupKeys.entrySet()) {
@@ -110,7 +111,7 @@ public class RawKVClient implements AutoCloseable {
           batches,
           entry.getKey(),
           entry.getValue(),
-          entry.getValue().stream().map(keysToValues::get).collect(Collectors.toList()),
+          entry.getValue().stream().map(kvPairs::get).collect(Collectors.toList()),
           RAW_BATCH_PUT_SIZE);
     }
     sendBatchPut(backOffer, batches);
@@ -238,7 +239,7 @@ public class RawKVClient implements AutoCloseable {
    * @param keys keys
    * @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<>();
     TiRegion lastRegion = null;
     for (ByteString key : keys) {
diff --git a/src/test/java/org/tikv/raw/RawKVClientTest.java b/src/test/java/org/tikv/raw/RawKVClientTest.java
index 312cedc6c9..66588474db 100644
--- a/src/test/java/org/tikv/raw/RawKVClientTest.java
+++ b/src/test/java/org/tikv/raw/RawKVClientTest.java
@@ -1,6 +1,5 @@
 package org.tikv.raw;
 
-import static org.tikv.raw.RawKVClient.mapKeysToValues;
 
 import com.google.protobuf.ByteString;
 import java.util.*;
@@ -259,15 +258,13 @@ public class RawKVClientTest {
         int i = cnt;
         completionService.submit(
             () -> {
-              List<ByteString> keyList = new ArrayList<>();
-              List<ByteString> valueList = new ArrayList<>();
+              Map<ByteString, ByteString> kvPairs = new HashMap<>();
               for (int j = 0; j < base; j++) {
                 int num = i * base + j;
                 ByteString key = orderedKeys.get(num), value = values.get(num);
-                keyList.add(key);
-                valueList.add(value);
+                kvPairs.put(key, value);
               }
-              client.batchPut(keyList, valueList);
+              client.batchPut(kvPairs);
               return null;
             });
       }
@@ -282,15 +279,13 @@ public class RawKVClientTest {
               + " put="
               + rawKeys().size());
     } else {
-      List<ByteString> keyList = new ArrayList<>();
-      List<ByteString> valueList = new ArrayList<>();
+      Map<ByteString, ByteString> kvPairs = new HashMap<>();
       for (int i = 0; i < putCases; i++) {
         ByteString key = randomKeys.get(i), value = values.get(r.nextInt(KEY_POOL_SIZE));
         data.put(key, value);
-        keyList.add(key);
-        valueList.add(value);
+        kvPairs.put(key, value);
       }
-      checkBatchPut(keyList, valueList);
+      checkBatchPut(kvPairs);
     }
   }
 
@@ -402,11 +397,10 @@ public class RawKVClientTest {
     assert client.get(key).equals(value);
   }
 
-  private void checkBatchPut(List<ByteString> keyList, List<ByteString> valueList) {
-    client.batchPut(keyList, valueList);
-    Map<ByteString, ByteString> keysToValues = mapKeysToValues(keyList, valueList);
-    for (ByteString key : keyList) {
-      assert client.get(key).equals(keysToValues.get(key));
+  private void checkBatchPut(Map<ByteString, ByteString> kvPairs) {
+    client.batchPut(kvPairs);
+    for (Map.Entry<ByteString, ByteString> kvPair : kvPairs.entrySet()) {
+      assert client.get(kvPair.getKey()).equals(kvPair.getValue());
     }
   }