[Java] Local and distributed ref counting in Java (#9371)

java/api/src/main/java/io/ray/api/ObjectRef.java

@@ -1,7 +1,5 @@
 package io.ray.api;
 
-import io.ray.api.id.ObjectId;
-
 /**
  * Represents a reference to an object in the object store.
  * @param <T> The object type.
@@ -14,15 +12,5 @@ public interface ObjectRef<T> {
    */
   T get();
 
-  /**
-   * Get the object id.
-   */
-  ObjectId getId();
-
-  /**
-   * Get the Object type.
-   */
-  Class<T> getType();
-
 }
 

java/api/src/main/java/io/ray/api/Ray.java

@@ -1,13 +1,11 @@
 package io.ray.api;
 
-import io.ray.api.id.ObjectId;
 import io.ray.api.id.UniqueId;
 import io.ray.api.placementgroup.PlacementGroup;
 import io.ray.api.placementgroup.PlacementStrategy;
 import io.ray.api.runtime.RayRuntime;
 import io.ray.api.runtime.RayRuntimeFactory;
 import io.ray.api.runtimecontext.RuntimeContext;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -67,25 +65,13 @@ public final class Ray extends RayCall {
   }
 
   /**
-   * Get an object by id from the object store.
+   * Get an object by `ObjectRef` from the object store.
    *
-   * @param objectId The ID of the object to get.
-   * @param objectType The type of the object to get.
+   * @param objectRef The reference of the object to get.
    * @return The Java object.
    */
-  public static <T> T get(ObjectId objectId, Class<T> objectType) {
-    return runtime.get(objectId, objectType);
-  }
-
-  /**
-   * Get a list of objects by ids from the object store.
-   *
-   * @param objectIds The list of object IDs.
-   * @param objectType The type of object.
-   * @return A list of Java objects.
-   */
-  public static <T> List<T> get(List<ObjectId> objectIds, Class<T> objectType) {
-    return runtime.get(objectIds, objectType);
+  public static <T> T get(ObjectRef<T> objectRef) {
+    return runtime.get(objectRef);
   }
 
   /**
@@ -95,13 +81,7 @@ public final class Ray extends RayCall {
    * @return A list of Java objects.
    */
   public static <T> List<T> get(List<ObjectRef<T>> objectList) {
-    List<ObjectId> objectIds = new ArrayList<>();
-    Class<T> objectType = null;
-    for (ObjectRef<T> o : objectList) {
-      objectIds.add(o.getId());
-      objectType = o.getType();
-    }
-    return runtime.get(objectIds, objectType);
+    return runtime.get(objectList);
   }
 
   /**

java/api/src/main/java/io/ray/api/runtime/RayRuntime.java

@@ -10,7 +10,6 @@ import io.ray.api.function.PyActorMethod;
 import io.ray.api.function.PyFunction;
 import io.ray.api.function.RayFunc;
 import io.ray.api.id.ActorId;
-import io.ray.api.id.ObjectId;
 import io.ray.api.id.UniqueId;
 import io.ray.api.options.ActorCreationOptions;
 import io.ray.api.options.CallOptions;
@@ -43,20 +42,18 @@ public interface RayRuntime {
   /**
    * Get an object from the object store.
    *
-   * @param objectId The ID of the object to get.
-   * @param objectType The type of the object to get.
+   * @param objectRef The reference of the object to get.
    * @return The Java object.
    */
-  <T> T get(ObjectId objectId, Class<T> objectType);
+  <T> T get(ObjectRef<T> objectRef);
 
   /**
    * Get a list of objects from the object store.
    *
-   * @param objectIds The list of object IDs.
-   * @param objectType The type of object.
+   * @param objectRefs The list of object references.
    * @return A list of Java objects.
    */
-  <T> List<T> get(List<ObjectId> objectIds, Class<T> objectType);
+  <T> List<T> get(List<ObjectRef<T>> objectRefs);
 
   /**
    * Wait for a list of RayObjects to be locally available, until specified number of objects are
@@ -72,11 +69,11 @@ public interface RayRuntime {
   /**
    * Free a list of objects from Plasma Store.
    *
-   * @param objectIds The object ids to free.
+   * @param objectRefs The object references to free.
    * @param localOnly Whether only free objects for local object store or not.
    * @param deleteCreatingTasks Whether also delete objects' creating tasks from GCS.
    */
-  void free(List<ObjectId> objectIds, boolean localOnly, boolean deleteCreatingTasks);
+  void free(List<ObjectRef<?>> objectRefs, boolean localOnly, boolean deleteCreatingTasks);
 
   /**
    * Set the resource for the specific node.