code: def convert_predict_rdd_to_xshard(data, prediction_rdd):
      import numpy as np
      from zoo.orca.data import SparkXShards
  
      def group_index(iter):
          for data in iter:
              size = get_size(data["x"])
              for i in range(size):
                  yield size
  
      def transform_predict(predictions):
          # list of np array
          if isinstance(predictions[0], list):
              predictions = np.array(predictions).T.tolist()
              result = [np.array(predict) for predict in predictions]
              return result
          # np array
          else:
              return np.array(predictions)
  
      def group(iter):
          this_index = 0
          buffer = []
          this_count = None
          for count, pred in iter:
              if this_index == 0:
                  this_count = count
              if this_index < this_count:
                  buffer.append(pred)
                  this_index += 1
              if this_index == this_count:
                  yield transform_predict(buffer)
                  buffer.clear()
                  this_index = 0
  
      def add_pred(shard_pred):
          shard, pred = shard_pred
          shard["prediction"] = pred
          return shard
  
      indexed_rdd = data.rdd.mapPartitions(group_index)
      grouped_pred = indexed_rdd.zip(prediction_rdd).mapPartitions(
          group
      )
      result_rdd = data.rdd.zip(grouped_pred).map(add_pred)
      return SparkXShards(result_rdd)
  
file_path: analytics-zoo/pyzoo/zoo/orca/learn/utils.py
imports: from zoo.common.utils import get_file_list
  from zoo.orca.data import SparkXShards
  from zoo.orca.data.utils import get_size
  from zoo.util.utils import convert_row_to_numpy
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: convert_predict_rdd_to_xshard(data, prediction_rdd) and returns: SparkXShards, shard, result, np