com.clickzetta.platform.test.BaseIgsWorkerV2 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.test;

import com.google.common.annotations.VisibleForTesting;
import cz.proto.ingestion.v2.IngestionV2;
import cz.proto.ingestion.v2.IngestionWorkerServiceGrpc;
import io.grpc.stub.StreamObserver;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowStreamReader;
import org.apache.arrow.vector.validate.ValidateVectorDataVisitor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@VisibleForTesting
public class BaseIgsWorkerV2 extends IngestionWorkerServiceGrpc.IngestionWorkerServiceImplBase {

  private static final Logger LOG = LoggerFactory.getLogger(BaseIgsWorkerV2.class);

  public static final Map>>> operatorMaps = new ConcurrentHashMap<>();

  public List>> getAllOperators() {
    List>> result = new ArrayList<>();
    for (List>> list : operatorMaps.values()) {
      result.addAll(list);
    }
    return result;
  }

  public void cleanStatus() {
    operatorMaps.clear();
  }

  public void decodeRequest(IngestionV2.MutateRequest request,
                            Map>>> operatorMaps) throws IOException {
    List>> arrowRows = new ArrayList<>();
    IngestionV2.DataBlock dataBlock = request.getDataBlock();
    for (int index = 0; index < dataBlock.getNumRows(); index++) {
      IngestionV2.OperationType operationType = null;
      if (dataBlock.hasBlockOpType()) {
        operationType = dataBlock.getBlockOpType();
      } else {
        operationType = dataBlock.getRowOpTypeList().getOpTypes(index);
      }
      arrowRows.add(new Tuple2<>(operationType, new ArrayList<>()));
    }

    try (BufferAllocator bufferAllocator = new RootAllocator();
         ArrowStreamReader streamReader = new ArrowStreamReader(dataBlock.getArrowPayload().newInput(), bufferAllocator)) {
      ValidateVectorDataVisitor dataVisitor = new ValidateVectorDataVisitor();
      while (streamReader.loadNextBatch()) {
        VectorSchemaRoot root = streamReader.getVectorSchemaRoot();
        List vectors = root.getFieldVectors();
        for (FieldVector vector : vectors) {
          vector.accept(dataVisitor, null);
        }
        for (int rowIdx = 0; rowIdx < root.getRowCount(); rowIdx++) {
          for (FieldVector vector : vectors) {
            arrowRows.get(rowIdx)._2.add(vector.getObject(rowIdx));
          }
        }
      }
    }
    List>> list =
        operatorMaps.computeIfAbsent(request.getBatchId(), s -> new ArrayList<>());
    list.addAll(arrowRows);
  }

  @Override
  public StreamObserver mutate(StreamObserver responseObserver) {
    return new StreamObserver() {
      @Override
      public void onNext(IngestionV2.MutateRequest request) {
        IngestionV2.MutateResponse response = IngestionV2.MutateResponse.newBuilder()
            .setBatchId(request.getBatchId())
            .setStatus(IngestionV2.ResponseStatus.newBuilder().setCode(IngestionV2.Code.SUCCESS).build())
            .build();
        try {
          decodeRequest(request, operatorMaps);
        } catch (IOException e) {
          throw new RuntimeException(e);
        }
        responseObserver.onNext(response);
      }

      @Override
      public void onError(Throwable throwable) {
        responseObserver.onError(throwable);
        LOG.error("BaseIgsWorkerV2 responseObserver onError.", throwable);
      }

      @Override
      public void onCompleted() {
        responseObserver.onCompleted();
      }
    };
  }

  @Override
  public StreamObserver multiMutate(StreamObserver responseObserver) {
    return new StreamObserver() {
      @Override
      public void onNext(IngestionV2.MultiMutateRequest request) {
        IngestionV2.MultiMutateResponse response = IngestionV2.MultiMutateResponse.newBuilder()
            .setBatchId(request.getBatchId())
            .setStatus(IngestionV2.ResponseStatus.newBuilder().setCode(IngestionV2.Code.SUCCESS).build())
            .build();
        for (IngestionV2.MutateRequest mutateRequest : request.getMutateRequestsList()) {
          try {
            decodeRequest(mutateRequest, operatorMaps);
          } catch (IOException e) {
            throw new RuntimeException(e);
          }
        }
        responseObserver.onNext(response);
      }

      @Override
      public void onError(Throwable throwable) {
        responseObserver.onError(throwable);
      }

      @Override
      public void onCompleted() {
        responseObserver.onCompleted();
      }
    };
  }
}