com.clickzetta.platform.flusher.ArrowFlushTask Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of clickzetta-java Show documentation
The java SDK for clickzetta's Lakehouse
There is a newer version: 2.0.0
package com.clickzetta.platform.flusher;

import com.clickzetta.platform.arrow.ArrowRecordBatchWriter;
import com.clickzetta.platform.arrow.ArrowTable;
import com.clickzetta.platform.client.CZIgsContext;
import com.clickzetta.platform.client.RpcRequestCallback;
import com.clickzetta.platform.client.api.ArrowRow;
import com.clickzetta.platform.client.api.ClientContext;
import com.clickzetta.platform.client.api.Listener;
import com.clickzetta.platform.connection.ChannelManager;
import com.clickzetta.platform.metrics.NettyMemoryMetrics;
import com.clickzetta.platform.operator.WriteOperation;
import com.google.common.base.Preconditions;
import com.google.protobuf.ByteStringWrap;
import cz.proto.ingestion.Ingestion;
import cz.proto.ingestion.v2.IngestionV2;
import scala.Tuple2;

import java.util.List;
import java.util.function.Supplier;

public class ArrowFlushTask extends AbstractTask {
  private long batchId;
  private ClientContext context;
  private Buffer buffer;
  private String serverToken;
  private Supplier> channelDataSupplier;
  private RpcRequestCallback rpcRequestCallback;
  private Listener listener;

  private boolean pooledAllocatorSupport;

  public ArrowFlushTask(
      long batchId,
      ClientContext context,
      Buffer buffer,
      String serverToken,
      Supplier> channelDataSupplier,
      RpcRequestCallback rpcResponseCallback,
      Listener listener) {
    this.batchId = batchId;
    this.context = context;
    this.buffer = buffer;
    this.serverToken = serverToken;
    this.channelDataSupplier = channelDataSupplier;
    this.rpcRequestCallback = rpcResponseCallback;
    this.listener = listener;
    this.pooledAllocatorSupport = ((CZIgsContext) context).isPooledAllocatorSupport();
  }

  @Override
  public void callPrepare() throws Exception {
    NettyMemoryMetrics.getInstance(false);
  }

  @Override
  public void callInternal() throws Exception {
    if (!buffer.isEmpty()) {
      ArrowTable table = (ArrowTable) buffer.getOperations().get(0).getTable();
      int numRows = buffer.getCurrentLines();
      List operationList = buffer.getOperations();

      List operationTypes;
      byte[] isSetBitMap;
      byte[] arrow_payload;
      ArrowRecordBatchWriter recordBatchWriter = new ArrowRecordBatchWriter(table, pooledAllocatorSupport, buffer.getCurrentLines());
      // encode arrow data format.
      try {
        for (WriteOperation writeOperation : operationList) {
          recordBatchWriter.write((ArrowRow) writeOperation);
        }
        recordBatchWriter.finish();

        // get operationTypes first. (not copy list. be careful with recordBatchWriter.reset)
        operationTypes = recordBatchWriter.getOperationTypes();
        isSetBitMap = recordBatchWriter.encodeIsSetBitMaps();
        arrow_payload = recordBatchWriter.encodeArrowRow();
      } finally {
        recordBatchWriter.close();
        recordBatchWriter = null;
      }

      IngestionV2.Account account = IngestionV2.Account.newBuilder()
          .setUserIdent(IngestionV2.UserIdentifier.newBuilder()
              .setInstanceId(context.instanceId())
              .setWorkspace(context.workspace()).build())
          .build();

      IngestionV2.MutateRequest.Builder requestBuilder = IngestionV2.MutateRequest.newBuilder()
          .setBatchId(batchId)
          .setWriteTimestamp(System.currentTimeMillis())
          .setTableIdent(IngestionV2.TableIdentifier.newBuilder()
              .setInstanceId(context.instanceId())
              .setWorkspace(context.workspace())
              .setSchemaName(table.getSchemaName())
              .setTableName(table.getTableName()).build())
          .setAccount(account);

      if (serverToken != null) {
        requestBuilder.setServerToken(serverToken);
      }

      IngestionV2.DataBlock.Builder dataBlockBuilder = IngestionV2.DataBlock.newBuilder()
          .setArrowPayload(ByteStringWrap.wrap(arrow_payload))
          .setIsSetBitmapsPayload(ByteStringWrap.wrap(isSetBitMap))
          .setNumRows(numRows);
      Tuple2 tuple2 = buildOpTypeInfo(operationTypes, table.getIgsTableType());
      if (tuple2._1 == IngestionV2.OperationType.class) {
        dataBlockBuilder.setBlockOpType((IngestionV2.OperationType) tuple2._2);
      } else {
        dataBlockBuilder.setRowOpTypeList((IngestionV2.OperationTypeList) tuple2._2);
      }
      requestBuilder.setDataBlock(dataBlockBuilder.build());

      if (context.authentication()) {
        IngestionV2.UserIdentifier.Builder userIdentifierBuilder = IngestionV2.UserIdentifier.newBuilder()
            .setInstanceId(context.instanceId())
            .setWorkspace(context.workspace())
            .setUserName(context.userName());
        if (context.userId() != null) {
          userIdentifierBuilder.setUserId(context.userId());
        }
        IngestionV2.Account.Builder accountBuilder = account.toBuilder();
        accountBuilder.setUserIdent(userIdentifierBuilder.build());
        accountBuilder.setToken(context.token());
        requestBuilder.setAccount(accountBuilder.build());
      }
      IngestionV2.MutateRequest request = requestBuilder.build();
      ChannelManager.ChannelData channelData = channelDataSupplier.get();
      rpcRequestCallback.setTargetHost(channelData.hostPort.toString());
      rpcRequestCallback.onSuccess(request, future);
      try {
        synchronized (channelData.streamObserver) {
          channelData.streamObserver.onNext(request);
        }
      } catch (Throwable t) {
        // roll back rpcRequestCallback onSuccess.
        rpcRequestCallback.onFailure(request, future, t);
      }
    } else {
      future.complete(true);
    }
    listener.onApplyRpcSend();
  }

  private  Tuple2 buildOpTypeInfo(List operationTypeList,
                                               Ingestion.IGSTableType igsTableType) {
    if (igsTableType != Ingestion.IGSTableType.ACID) {
      Preconditions.checkArgument(operationTypeList.size() == 1 && operationTypeList.get(0) == IngestionV2.OperationType.INSERT,
          "Common & ClusterTable only support Insert Operation.");
      return new Tuple2<>(IngestionV2.OperationType.class, (T) operationTypeList.get(0));
    } else {
      return new Tuple2<>(IngestionV2.OperationTypeList.class,
          (T) IngestionV2.OperationTypeList.newBuilder().addAllOpTypes(operationTypeList).build());
    }
  }

  @Override
  public long getId() {
    return batchId;
  }

  public Buffer getBuffer() {
    return buffer;
  }
}