code: def __init__(
      self,
      meta_data,
      object_store_address,
      workers_per_node=1,
      batch_size=1,
  ):
      import pyarrow.plasma as plasma
  
      self.client = plasma.connect(object_store_address)
      print("Connected to plasma")
  
      # All the subpartitions on this node
      all_data = [
          subpartition
          for subpartition in meta_data
          if subpartition[4] == get_node_ip()
      ]
      rank = int(os.environ.get("PMI_RANK", 0))
      print("Global rank: ", rank)
      # rank = int(os.environ.get("PMIX_RANK", 0))  # For OpenMPI
      local_rank = rank % workers_per_node
      print("Local rank: ", local_rank)
      data_splits = list(
          chunks(all_data, len(all_data) // workers_per_node)
      )
      worker_data = data_splits[local_rank]
      if len(data_splits) == (
          workers_per_node + 1
      ):  # Can't evenly split among workers
          remain_data = data_splits[-1]
          if local_rank < len(remain_data):
              worker_data += [remain_data[local_rank]]
      self.object_ids = [
          subpartition[3] for subpartition in worker_data
      ]
      self.sizes = [subpartition[2] for subpartition in worker_data]
      print("Data size for worker: ", sum(self.sizes))
      self.batch_size = batch_size
      offsets = []
      for i in self.sizes:
          if len(offsets) == 0:
              offsets.append(i)
          else:
              offsets.append(offsets[-1] + i)
      self.offsets = offsets
      self.current_index = 0  # Current index for object_id; data loaded
      self.load_from_plasma(self.current_index)
  
imports: import os
  import subprocess
  import types
  import cloudpickle
  from pyspark.rdd import RDD
  from pyspark.sql import DataFrame
  from torch.utils.data import Dataset, DataLoader
  from bigdl.dllib.utils.log4Error import *
  from bigdl.dllib.utils.utils import get_node_ip
  from bigdl.orca.learn.mpi.mpi_runner import MPIRunner
  from bigdl.orca.learn.mpi.utils import *
label_desc: Use image node_id 1 to create a new PlasmaNDArrayDataset object from inherited base classes: Dataset with example: obj = PlasmaNDArrayDataset(meta_data, object_store_address, workers_per_node, batch_size)