eliminate usage of deserialize_and_register_object_ref

Author: pang-wu

core/raydp-main/src/main/scala/org/apache/spark/sql/raydp/ObjectStoreWriter.scala

@@ -67,42 +67,35 @@ class ObjectStoreWriter(@transient val df: DataFrame) extends Serializable {
       queue: ObjectRefHolder.Queue,
       ownerName: String): RecordBatch = {
 
-    // NOTE: We intentionally do NOT pass an owner argument to Ray.put anymore.
-    //
-    // - When ownerName is empty, route the put via the JVM RayAppMaster actor.
-    // - When ownerName is set to a Python actor name (e.g. RayDPSparkMaster),
-    //   invoke that Python actor's put_data(data) method via Ray cross-language
-    //   calls so that the Python actor becomes the owner of the created object.
-    val objectRef: ObjectRef[_] =
-      if (ownerName == "") {
-        Ray.put(data)
-      } else {
-        // Ray.getActor(String) is a raw Java Optional in Ray's Java API.
-        // If we don't cast it to an explicit reference type here, Scala may infer
-        // Optional[Nothing] and insert an invalid cast at runtime.
-        val opt = Ray.getActor(ownerName).asInstanceOf[Optional[AnyRef]]
-        if (!opt.isPresent) {
-          throw new RayDPException(s"Actor $ownerName not found when putting dataset block.")
-        }
-        val handleAny: AnyRef = opt.get()
-        if (!handleAny.isInstanceOf[PyActorHandle]) {
-          throw new RayDPException(
-            s"Actor $ownerName is not a Python actor; cannot invoke put_data."
-          )
-        }
-        val pyHandle = handleAny.asInstanceOf[PyActorHandle]
-        val method = PyActorMethod.of("put_data", classOf[AnyRef])
-        val refOfRef = pyHandle.task(method, data).remote()
-        refOfRef
-      }
+    // Owner-transfer only implementation:
+    // - ownerName must always be provided (non-empty) and refer to a Python actor.
+    // - JVM never creates/handles Ray ObjectRefs for the dataset blocks.
+    // - JVM returns only a per-batch key encoded in RecordBatch.objectId (bytes),
+    //   and Python will fetch the real ObjectRefs from the owner actor by key.
+
+    if (ownerName == null || ownerName.isEmpty) {
+      throw new RayDPException("ownerName must be set for Spark->Ray conversion.")
+    }
+
+    val opt = Ray.getActor(ownerName).asInstanceOf[Optional[AnyRef]]
+    if (!opt.isPresent) {
+      throw new RayDPException(s"Actor $ownerName not found when putting dataset block.")
+    }
+    val handleAny: AnyRef = opt.get()
+    if (!handleAny.isInstanceOf[PyActorHandle]) {
+      throw new RayDPException(s"Actor $ownerName is not a Python actor; cannot invoke put_data.")
+    }
+    val pyHandle = handleAny.asInstanceOf[PyActorHandle]
+    val batchKey = UUID.randomUUID().toString
+
+    // put_data(batchKey, arrowBytes) -> boolean ack
+    val method = PyActorMethod.of("put_data", classOf[java.lang.Boolean])
+    val args: Array[AnyRef] = Array(batchKey, data.asInstanceOf[AnyRef])
+    new PyActorTaskCaller(pyHandle, method, args).remote().get()
 
-    // add the objectRef to the objectRefHolder to avoid reference GC
-    queue.add(objectRef)
-    val objectRefImpl = RayDPUtils.convert(objectRef)
-    val objectId = objectRefImpl.getId
-    val runtime = Ray.internal.asInstanceOf[AbstractRayRuntime]
-    val addressInfo = runtime.getObjectStore.getOwnershipInfo(objectId)
-    RecordBatch(addressInfo, objectId.getBytes, numRecords)
+    // ownerAddress/objectId here are not Ray's object metadata; objectId encodes the key.
+    // Python side will treat objectId as UTF-8 key bytes.
+    RecordBatch(Array.emptyByteArray, batchKey.getBytes("UTF-8"), numRecords)
   }
 
   /**

python/raydp/spark/dataset.py

@@ -110,39 +110,26 @@ def raydp_master_set_reference_as_state(
         raydp_master_set_reference_as_state)
 
 
-@client_mode_wrap
-def _register_objects(records):
-    worker = ray.worker.global_worker
-    blocks: List[ray.ObjectRef] = []
-    block_sizes: List[int] = []
-    for obj_id, owner, num_record in records:
-        object_ref = ray.ObjectRef(obj_id)
-        # Register the ownership of the ObjectRef
-        worker.core_worker.deserialize_and_register_object_ref(
-            object_ref.binary(), ray.ObjectRef.nil(), owner, "")
-        blocks.append(object_ref)
-        block_sizes.append(num_record)
-    return blocks, block_sizes
-
 def _save_spark_df_to_object_store(df: sql.DataFrame, use_batch: bool = True,
                                    owner: Union[PartitionObjectsOwner, None] = None):
     # call java function from python
     jvm = df.sql_ctx.sparkSession.sparkContext._jvm
     jdf = df._jdf
     object_store_writer = jvm.org.apache.spark.sql.raydp.ObjectStoreWriter(jdf)
-    actor_owner_name = ""
-    if owner is not None:
-        actor_owner_name = owner.actor_name
+    if owner is None:
+        # Default to RayDPSparkMaster as the owner if not specified.
+        owner = get_raydp_master_owner(df.sql_ctx.sparkSession)
+    actor_owner_name = owner.actor_name
     records = object_store_writer.save(use_batch, actor_owner_name)
 
-    record_tuples = [(record.objectId(), record.ownerAddress(), record.numRecords())
-                     for record in records]
-    blocks, block_sizes = _register_objects(record_tuples)
-
-    if owner is not None:
-        actor_owner = ray.get_actor(actor_owner_name)
-        ray.get(owner.set_reference_as_state(actor_owner, blocks))
-
+    # Owner-transfer-only path:
+    # JVM returns List[RecordBatch] where record.objectId() contains UTF-8 bytes of batch_key.
+    # Fetch actual ObjectRefs from the owner actor by key.
+    data_owner_actor = ray.get_actor(actor_owner_name)
+    batch_keys = [bytes(record.objectId()).decode("utf-8") for record in records]
+    block_sizes = [record.numRecords() for record in records]
+    blocks = ray.get(data_owner_actor.get_block_refs.remote(batch_keys))
+    ray.get(owner.set_reference_as_state(data_owner_actor, blocks))
     return blocks, block_sizes
 
 def spark_dataframe_to_ray_dataset(df: sql.DataFrame,

python/raydp/spark/ray_cluster_master.py

@@ -57,12 +57,36 @@ class RayDPObjectOwnerMixin:
     objects, without using Ray's experimental `ray.put(_owner=...)` API.
     """
 
-    def put_data(self, data) -> "pa.Table":
-        """Put one serialized Arrow batch into the Ray object store."""
-        # data is Arrow IPC stream bytes written by ArrowStreamWriter
+    def _get_raydp_blocks_by_key(self):
+        blocks = getattr(self, "_raydp_blocks_by_key", None)
+        if blocks is None:
+            blocks = {}
+            setattr(self, "_raydp_blocks_by_key", blocks)
+        return blocks
+
+    def put_data(self, batch_key: str, data: bytes) -> bool:
+        """Create one Ray Dataset block owned by this actor.
+
+        Args:
+            batch_key: A per-batch application-level key generated by the JVM.
+            data: Arrow IPC stream bytes written by ArrowStreamWriter on Spark executors.
+
+        Returns:
+            True when the block has been created and stored.
+        """
         reader = pa.ipc.open_stream(pa.BufferReader(data))
         table = reader.read_all()
-        return table
+        ref = ray.put(table)
+        self._get_raydp_blocks_by_key()[batch_key] = ref
+        return True
+
+    def get_block_refs(self, batch_keys):
+        """Fetch (and remove) stored block refs for the given keys."""
+        blocks = self._get_raydp_blocks_by_key()
+        refs = []
+        for k in batch_keys:
+            refs.append(blocks.pop(k))
+        return refs