From 11bac252fd19c52fe92689b25ff6d8cc1304e47c Mon Sep 17 00:00:00 2001
From: iosmanthus <myosmanthustree@gmail.com>
Date: Wed, 9 Jun 2021 15:23:35 +0800
Subject: [PATCH] add batch scan keys

Signed-off-by: iosmanthus <myosmanthustree@gmail.com>
---
 src/main/java/org/tikv/raw/RawKVClient.java   | 43 +++++++++++++++++++
 .../java/org/tikv/raw/RawKVClientTest.java    | 18 ++++++++
 2 files changed, 61 insertions(+)

diff --git a/src/main/java/org/tikv/raw/RawKVClient.java b/src/main/java/org/tikv/raw/RawKVClient.java
index bf4feed4213..9e590e20f5b 100644
--- a/src/main/java/org/tikv/raw/RawKVClient.java
+++ b/src/main/java/org/tikv/raw/RawKVClient.java
@@ -345,6 +345,49 @@ public Long getKeyTTL(ByteString key) {
     }
   }
 
+  /**
+   * Create a new `batch scan` request with `keyOnly` option Once resolved this request will result
+   * in a set of scanners over the given keys.
+   *
+   * <p>WARNING: This method is experimental. The `each_limit` parameter does not work as expected.
+   * It does not limit the number of results returned of each range, instead it limits the number of
+   * results in each region of each range. As a result, you may get more than each_limit key-value
+   * pairs for each range. But you should not miss any entries.
+   *
+   * @param ranges a list of ranges
+   * @return a set of scanners for keys over the given keys.
+   */
+  public List<List<ByteString>> batchScanKeys(
+      List<Pair<ByteString, ByteString>> ranges, int eachLimit) {
+    return batchScan(
+            ranges
+                .stream()
+                .map(
+                    range ->
+                        ScanOption.newBuilder()
+                            .setStartKey(range.first)
+                            .setEndKey(range.second)
+                            .setLimit(eachLimit)
+                            .setKeyOnly(true)
+                            .build())
+                .collect(Collectors.toList()))
+        .stream()
+        .map(kvs -> kvs.stream().map(kv -> kv.getKey()).collect(Collectors.toList()))
+        .collect(Collectors.toList());
+  }
+
+  /**
+   * Create a new `batch scan` request. Once resolved this request will result in a set of scanners
+   * over the given keys.
+   *
+   * <p>WARNING: This method is experimental. The `each_limit` parameter does not work as expected.
+   * It does not limit the number of results returned of each range, instead it limits the number of
+   * results in each region of each range. As a result, you may get more than each_limit key-value
+   * pairs for each range. But you should not miss any entries.
+   *
+   * @param ranges a list of `ScanOption` for each range
+   * @return a set of scanners over the given keys.
+   */
   public List<List<KvPair>> batchScan(List<ScanOption> ranges) {
     String label = "client_raw_batch_scan";
     Histogram.Timer requestTimer = RAW_REQUEST_LATENCY.labels(label).startTimer();
diff --git a/src/test/java/org/tikv/raw/RawKVClientTest.java b/src/test/java/org/tikv/raw/RawKVClientTest.java
index cda50aa830c..968a77f8a94 100644
--- a/src/test/java/org/tikv/raw/RawKVClientTest.java
+++ b/src/test/java/org/tikv/raw/RawKVClientTest.java
@@ -18,6 +18,7 @@
 import org.tikv.common.exception.TiKVException;
 import org.tikv.common.key.Key;
 import org.tikv.common.util.FastByteComparisons;
+import org.tikv.common.util.Pair;
 import org.tikv.common.util.ScanOption;
 import org.tikv.kvproto.Kvrpcpb;
 
@@ -644,6 +645,11 @@ private void rawBatchScanTest(int scanCases, boolean benchmark) {
         scanOptions.add(scanOption);
       }
       checkBatchScan(scanOptions);
+      checkBatchScanKeys(
+          scanOptions
+              .stream()
+              .map(scanOption -> Pair.create(scanOption.getStartKey(), scanOption.getEndKey()))
+              .collect(Collectors.toList()));
     }
   }
 
@@ -801,6 +807,7 @@ private void checkScan(
   }
 
   private void checkBatchScan(List<ScanOption> scanOptions) {
+    logger.info("checking batch scan");
     List<List<Kvrpcpb.KvPair>> result = client.batchScan(scanOptions);
     int i = 0;
     for (ScanOption scanOption : scanOptions) {
@@ -820,6 +827,17 @@ private void checkBatchScan(List<ScanOption> scanOptions) {
     }
   }
 
+  private void checkBatchScanKeys(List<Pair<ByteString, ByteString>> ranges) {
+    logger.info("checking batch scan keys");
+    List<List<ByteString>> result = client.batchScanKeys(ranges, limit);
+    for (int i = 0; i < ranges.size(); i++) {
+      Pair<ByteString, ByteString> range = ranges.get(i);
+      List<ByteString> partialResult =
+          new ArrayList<>(data.subMap(range.first, range.second).keySet());
+      assert result.get(i).equals(partialResult);
+    }
+  }
+
   private void checkDelete(ByteString key) {
     client.delete(key);
     checkEmpty(key);