From a781b103da1b2819c966fd098c941c037ce4cc20 Mon Sep 17 00:00:00 2001
From: Jay Zhuang <zjay@fb.com>
Date: Wed, 31 Mar 2021 09:45:47 -0700
Subject: [PATCH] Fix getApproximateMemTableStats() return type (#8098)

Summary:
Which should return 2 long instead of an array.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/8098

Reviewed By: mrambacher

Differential Revision: D27308741

Pulled By: jay-zhuang

fbshipit-source-id: 44beea2bd28cf6779b048bebc98f2426fe95e25c
---
 java/rocksjni/rocksjni.cc                     |  5 ++---
 .../test/java/org/rocksdb/RocksDBTest.java    | 20 +++++++++++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/java/rocksjni/rocksjni.cc b/java/rocksjni/rocksjni.cc
index 7dc0289e8..054c0fc89 100644
--- a/java/rocksjni/rocksjni.cc
+++ b/java/rocksjni/rocksjni.cc
@@ -2428,14 +2428,13 @@ jlongArray Java_org_rocksdb_RocksDB_getApproximateMemTableStats(
       static_cast<jlong>(count),
       static_cast<jlong>(sizes)};
 
-  const jsize jcount = static_cast<jsize>(count);
-  jlongArray jsizes = env->NewLongArray(jcount);
+  jlongArray jsizes = env->NewLongArray(2);
   if (jsizes == nullptr) {
     // exception thrown: OutOfMemoryError
     return nullptr;
   }
 
-  env->SetLongArrayRegion(jsizes, 0, jcount, results);
+  env->SetLongArrayRegion(jsizes, 0, 2, results);
   if (env->ExceptionCheck()) {
     // exception thrown: ArrayIndexOutOfBoundsException
     env->DeleteLocalRef(jsizes);
diff --git a/java/src/test/java/org/rocksdb/RocksDBTest.java b/java/src/test/java/org/rocksdb/RocksDBTest.java
index 46f6414c8..643442352 100644
--- a/java/src/test/java/org/rocksdb/RocksDBTest.java
+++ b/java/src/test/java/org/rocksdb/RocksDBTest.java
@@ -1271,6 +1271,26 @@ public class RocksDBTest {
     }
   }
 
+  @Test
+  public void getApproximateMemTableStatsSingleKey() throws RocksDBException {
+    final byte key1[] = "key1".getBytes(UTF_8);
+    final byte key2[] = "key2".getBytes(UTF_8);
+    final byte key3[] = "key3".getBytes(UTF_8);
+    try (final Options options = new Options().setCreateIfMissing(true)) {
+      final String dbPath = dbFolder.getRoot().getAbsolutePath();
+      try (final RocksDB db = RocksDB.open(options, dbPath)) {
+        db.put(key1, key1);
+
+        final RocksDB.CountAndSize stats =
+            db.getApproximateMemTableStats(new Range(new Slice(key1), new Slice(key3)));
+
+        assertThat(stats).isNotNull();
+        assertThat(stats.count).isEqualTo(1);
+        assertThat(stats.size).isGreaterThan(1);
+      }
+    }
+  }
+
   @Ignore("TODO(AR) re-enable when ready!")
   @Test
   public void compactFiles() throws RocksDBException {