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.
+ *
+ *
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> batchScanKeys(
+ List> 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.
+ *
+ * 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> batchScan(List 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 scanOptions) {
+ logger.info("checking batch scan");
List> result = client.batchScan(scanOptions);
int i = 0;
for (ScanOption scanOption : scanOptions) {
@@ -820,6 +827,17 @@ private void checkBatchScan(List scanOptions) {
}
}
+ private void checkBatchScanKeys(List> ranges) {
+ logger.info("checking batch scan keys");
+ List> result = client.batchScanKeys(ranges, limit);
+ for (int i = 0; i < ranges.size(); i++) {
+ Pair range = ranges.get(i);
+ List 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);