ARROW-8082: [Plasma] Add JNI list() interface

Currently Plasma list() method is not implemented in Java code. Implement PlasmaClientJNI.list() interface and add some UT. Closes #6588 from jikunshang/add_jni_list Lead-authored-by: Ji Kunshang <kunshang.ji@intel.com> Co-authored-by: kunshang <kunshang.ji@intel.com> Signed-off-by: Wes McKinney <wesm+git@apache.org>
apache · Apr 2, 2020 · fad2165 · fad2165
1 parent 9ad2cc4
commit fad2165
Show file tree

Hide file tree

Showing 6 changed files with 55 additions and 0 deletions.
diff --git a/cpp/src/plasma/lib/java/org_apache_arrow_plasma_PlasmaClientJNI.cc b/cpp/src/plasma/lib/java/org_apache_arrow_plasma_PlasmaClientJNI.cc
@@ -241,3 +241,23 @@ JNIEXPORT jlong JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_evict(
 
   return static_cast<jlong>(evicted_bytes);
 }
+
+JNIEXPORT jobjectArray JNICALL
+Java_org_apache_arrow_plasma_PlasmaClientJNI_list(JNIEnv* env, jclass cls, jlong conn) {
+  plasma::PlasmaClient* client = reinterpret_cast<plasma::PlasmaClient*>(conn);
+  plasma::ObjectTable objectTable;
+  throw_exception_if_not_OK(env, client->List(&objectTable));
+  jobjectArray ret =
+      env->NewObjectArray(objectTable.size(), env->FindClass("[B"), env->NewByteArray(1));
+  int i = 0;
+  for (const auto& id_entry_pair : objectTable) {
+    const plasma::ObjectID& id = id_entry_pair.first;
+    jbyteArray idByteArray = env->NewByteArray(OBJECT_ID_SIZE);
+    env->SetByteArrayRegion(idByteArray, 0, OBJECT_ID_SIZE,
+                            reinterpret_cast<jbyte*>(const_cast<uint8_t*>(id.data())));
+    env->SetObjectArrayElement(ret, i, idByteArray);
+    i++;
+  }
+
+  return ret;
+}
diff --git a/cpp/src/plasma/lib/java/org_apache_arrow_plasma_PlasmaClientJNI.h b/cpp/src/plasma/lib/java/org_apache_arrow_plasma_PlasmaClientJNI.h
@@ -126,6 +126,15 @@ JNIEXPORT jlong JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_evict(JNIEn
                                                                            jclass, jlong,
                                                                            jlong);
 
+/*
+ * Class:     org_apache_arrow_plasma_PlasmaClientJNI
+ * Method:    list
+ * Signature: (J)[[B
+ */
+JNIEXPORT jobjectArray JNICALL Java_org_apache_arrow_plasma_PlasmaClientJNI_list(JNIEnv*,
+                                                                                 jclass,
+                                                                                 jlong);
+
 #ifdef __cplusplus
 }
 #endif

diff --git a/java/plasma/src/main/java/org/apache/arrow/plasma/ObjectStoreLink.java b/java/plasma/src/main/java/org/apache/arrow/plasma/ObjectStoreLink.java
@@ -123,4 +123,9 @@ default byte[] get(byte[] objectId, int timeoutMs, boolean isMetadata) {
    * @param objectId used to identify an object.
    */
   boolean contains(byte[] objectId);
+
+  /**
+   * List all objects in the PlasmaStore.
+   */
+  List<byte[]> list();
 }
diff --git a/java/plasma/src/main/java/org/apache/arrow/plasma/PlasmaClient.java b/java/plasma/src/main/java/org/apache/arrow/plasma/PlasmaClient.java
@@ -19,6 +19,7 @@
 
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 import org.apache.arrow.plasma.exceptions.DuplicateObjectException;
@@ -118,6 +119,11 @@ public ByteBuffer getObjAsByteBuffer(byte[] objectId, int timeoutMs, boolean isM
     return bufs[0][isMetadata ? 1 : 0];
   }
 
+  @Override
+  public List<byte[]> list() {
+    return Arrays.asList(PlasmaClientJNI.list(conn));
+  }
+
   @Override
   public long evict(long numBytes) {
     return PlasmaClientJNI.evict(conn, numBytes);

diff --git a/java/plasma/src/main/java/org/apache/arrow/plasma/PlasmaClientJNI.java b/java/plasma/src/main/java/org/apache/arrow/plasma/PlasmaClientJNI.java
@@ -53,4 +53,5 @@ public static native byte[][] wait(long conn, byte[][] objectIds, int timeoutMs,
 
   public static native long evict(long conn, long numBytes);
 
+  public static native byte[][] list(long conn);
 }
diff --git a/java/plasma/src/test/java/org/apache/arrow/plasma/PlasmaClientTest.java b/java/plasma/src/test/java/org/apache/arrow/plasma/PlasmaClientTest.java
@@ -172,6 +172,20 @@ public void doTest() {
     Arrays.fill(value4, (byte)14);
     pLink.put(id4, value4, meta4);
 
+    List<byte[]> existIds = Arrays.asList(id1, id2, id3, id4);
+    List<byte[]> listIds = pLink.list();
+    assert listIds.size() == 4;
+    for (byte[] existId : existIds) {
+      boolean found = false;
+      for (byte[] listId : listIds) {
+        if (Arrays.equals(listId, existId)) {
+          found = true;
+        }
+      }
+      assert found;
+    }
+    System.out.println("Plasma java client list test success.");
+
     byte[] id5 = new byte[20];
     Arrays.fill(id5, (byte)5);
     byte[] value5 = new byte[20];