text: # Copyright 1999-2020 Alibaba Group Holding Ltd.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import numpy as np
  import pandas as pd
  
  from ... import opcodes as OperandDef
  from ...core import OutputType
  from ...core.operand.base import SchedulingHint
  from ...serialization.serializables import StringField
  from ...tensor.datastore.to_vineyard import resolve_vineyard_socket
  from ..operands import DataFrameOperand, DataFrameOperandMixin
  from ..utils import parse_index
  
  try:
      import vineyard
      from vineyard.data.dataframe import make_global_dataframe
      from vineyard.data.utils import to_json
  except ImportError:
      vineyard = None
  
  
  class DataFrameToVineyardChunk(DataFrameOperand, DataFrameOperandMixin):
      _op_type_ = OperandDef.DATAFRAME_STORE_VINEYARD_CHUNK
  
      # vineyard ipc socket
      vineyard_socket = StringField("vineyard_socket")
  
      def __init__(self, vineyard_socket=None, dtypes=None, **kw):
          super().__init__(
              vineyard_socket=vineyard_socket,
              _dtypes=dtypes,
              _output_types=[OutputType.dataframe],
              **kw
          )
  
      def __call__(self, df):
          return self.new_dataframe(
              [df],
              shape=(0, 0),
              dtypes=df.dtypes,
              index_value=df.index_value,
              columns_value=df.columns_value,
          )
  
      @classmethod
      def _process_out_chunks(cls, op, out_chunks):
          dtypes = pd.Series([np.dtype("O")], index=pd.Index([0]))
          merge_op = DataFrameToVinyardStoreMeta(
              vineyard_socket=op.vineyard_socket,
              chunk_shape=op.inputs[0].chunk_shape,
              shape=(1, 1),
              dtypes=dtypes,
          )
          return merge_op.new_chunks(
              out_chunks, shape=(1, 1), dtypes=dtypes, index=(0, 0)
          )
  
      @classmethod
      def tile(cls, op):
          out_chunks = []
          scheduling_hint = SchedulingHint(fuseable=False)
          dtypes = pd.Series([np.dtype("O")], index=pd.Index([0]))
          for idx, chunk in enumerate(op.inputs[0].chunks):
              chunk_op = op.copy().reset_key()
              chunk_op.scheduling_hint = scheduling_hint
              out_chunk = chunk_op.new_chunk(
                  [chunk],
                  shape=(1, 1),
                  dtypes=dtypes,
                  index_value=chunk.index_value,
                  columns_value=chunk.columns_value,
                  index=(idx, 0),
              )
              out_chunks.append(out_chunk)
          out_chunks = cls._process_out_chunks(op, out_chunks)
  
          in_df = op.inputs[0]
          new_op = op.copy().reset_key()
          return new_op.new_dataframes(
              op.inputs,
              shape=(len(out_chunks), 1),
              dtypes=dtypes,
              index_value=in_df.index_value,
              columns_value=in_df.columns_value,
              chunks=out_chunks,
              nsplits=((np.prod(op.inputs[0].chunk_shape),),),
          )
  
      @classmethod
      def execute(cls, ctx, op):
          if vineyard is None:
              raise RuntimeError("vineyard is not available")
  
          socket, needs_put = resolve_vineyard_socket(ctx, op)
          client = vineyard.connect(socket)
  
          # some op might be fused and executed twice on different workers
          if not needs_put:
              # might be fused
              try:  # pragma: no cover
                  meta = ctx.get_chunks_meta([op.inputs[0].key])[0]
                  df_id = vineyard.ObjectID(meta["object_ref"])
                  if not client.exists(df_id):
                      needs_put = True
              except KeyError:
                  needs_put = True
          if needs_put:
              df_id = client.put(
                  ctx[op.inputs[0].key], partition_index=op.inputs[0].index
              )
          else:  # pragma: no cover
              meta = client.get_meta(df_id)
              new_meta = vineyard.ObjectMeta()
              for k, v in meta.items():
                  if k not in ["id", "signature", "instance_id"]:
                      if isinstance(v, vineyard.ObjectMeta):
                          new_meta.add_member(k, v)
                      else:
                          new_meta[k] = v
              new_meta["partition_index_"] = to_json(op.inputs[0].index)
              df_id = client.create_metadata(new_meta).id
  
          client.persist(df_id)
          ctx[op.outputs[0].key] = pd.DataFrame({0: [df_id]})
  
  
  class DataFrameToVinyardStoreMeta(DataFrameOperand, DataFrameOperandMixin):
      _op_type_ = OperandDef.DATAFRAME_STORE_VINEYARD_META
  
      # vineyard ipc socket
      vineyard_socket = StringField("vineyard_socket")
  
      def __init__(self, vineyard_socket=None, dtypes=None, **kw):
          super().__init__(
              vineyard_socket=vineyard_socket,
              dtypes=dtypes,
              _output_types=[OutputType.dataframe],
              **kw
          )
  
      @classmethod
      def tile(cls, op):
          dtypes = pd.Series([np.dtype("O")], index=pd.Index([0]))
          chunk_op = op.copy().reset_key()
          out_chunk = chunk_op.new_chunk(
              op.inputs[0].chunks,
              shape=(1, 1),
              dtypes=dtypes,
              index_value=parse_index(pd.Index([-1])),
              columns_value=parse_index(pd.Index([0])),
              index=(0, 0),
          )
          new_op = op.copy().reset_key()
          return new_op.new_dataframes(
              op.inputs,
              shape=(1, 1),
              dtypes=dtypes,
              index_value=parse_index(pd.Index([0])),
              columns_value=parse_index(pd.Index([0])),
              chunks=[out_chunk],
              nsplits=((1,), (1,)),
          )
  
      @classmethod
      def execute(cls, ctx, op):
          if vineyard is None:
              raise RuntimeError("vineyard is not available")
  
          socket, _ = resolve_vineyard_socket(ctx, op)
          client = vineyard.connect(socket)
  
          # # store the result object id to execution context
          chunks = [ctx[chunk.key][0][0] for chunk in op.inputs]
          ctx[op.outputs[0].key] = pd.DataFrame(
              {0: [make_global_dataframe(client, chunks).id]}
          )
  
  
  def to_vineyard(df, vineyard_socket=None):
      op = DataFrameToVineyardChunk(vineyard_socket=vineyard_socket)
      return op(df)
  