com.clickzetta.platform.client.api.ArrowStream 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.client.api;

import com.clickzetta.platform.arrow.ArrowTable;
import com.clickzetta.platform.arrow.CZArrowSession;
import com.clickzetta.platform.client.CZClient;
import com.clickzetta.platform.client.pool.AbstractRowPoolSupport;
import com.clickzetta.platform.client.schemachange.EvolutionManager;
import com.clickzetta.platform.client.schemachange.EvolutionManagerImpl;
import com.clickzetta.platform.common.CZException;
import com.clickzetta.platform.common.Constant;
import com.clickzetta.platform.metrics.MetricBase;
import com.clickzetta.platform.operator.WriteOperation;
import cz.proto.ingestion.v2.IngestionV2;

import java.io.IOException;

/**
 * not thread safe.
 */
public final class ArrowStream extends AbstractRowPoolSupport implements Stream {

  ArrowTable table;
  CZClient client;
  CZArrowSession czArrowSession;

  private boolean complexTypeRreCheck;

  // A unique id for connected tablet server. Unavailable until the first successful mutate request.
  // It is used to prevent data loss on tablet restart and should not be changed during the session.
  // Optional for the 1st request, required for the following requests.
  volatile String serverToken;

  // Schema Change
  private EvolutionManager evolutionManager;

  public ArrowStream(CZClient client, ArrowTable table, Options options) {
    this.client = client;
    this.table = table;
    this.czArrowSession = CZArrowSession.build(client, this, options);
    this.serverToken = null;

    Object obj = options.getProperties().getOrDefault(Constant.ARROW_ROW_FORMAT_CHECK, true);
    this.complexTypeRreCheck = obj instanceof String ? Boolean.parseBoolean((String) obj) : (boolean) obj;

    // row pool support.
    initRowPool(options);

    // init schema change manager.
    this.evolutionManager = new EvolutionManagerImpl(options, client);
  }

  @Override
  public Row createInsertRow() {
    return getOrCreateRow(IngestionV2.OperationType.INSERT);
  }

  @Override
  public Row createInsertIgnoreRow() {
    return getOrCreateRow(IngestionV2.OperationType.INSERT_IGNORE);
  }

  @Override
  public Row createUpdateRow() {
    return getOrCreateRow(IngestionV2.OperationType.UPDATE);
  }

  @Override
  public Row createUpsertRow() {
    return getOrCreateRow(IngestionV2.OperationType.UPSERT);
  }

  @Override
  public Row createDeleteRow() {
    return getOrCreateRow(IngestionV2.OperationType.DELETE);
  }

  @Override
  public Row createDeleteIgnoreRow() {
    return getOrCreateRow(IngestionV2.OperationType.DELETE_IGNORE);
  }

  @Override
  protected Row buildEmptyRow() {
    return ArrowRow.buildReuseRow(complexTypeRreCheck);
  }

  private Row getOrCreateRow(IngestionV2.OperationType operationType) {
    WriteOperation operation = null;
    if (poolSupport()) {
      operation = getRowInPool().getWriteOperation();
      operation.resetRowMeta(table, operationType);
    } else {
      operation = new ArrowRow(table, operationType, complexTypeRreCheck);
    }
    return operation.getRow();
  }

  @Override
  public ProtocolType getProtocolType() {
    return ProtocolType.V2;
  }

  @Override
  public ArrowTable getTable() {
    return this.table;
  }

  @Override
  public EvolutionManager getEvolutionManager() {
    return this.evolutionManager;
  }

  public CZArrowSession getSession() {
    return this.czArrowSession;
  }

  @Override
  public void bindListener(Listener listener) {
    this.czArrowSession.bindListener(listener);
  }

  @Override
  public void apply(Row... rows) throws IOException {
    this.czArrowSession.apply(rows);
  }

  @Override
  public void flush() throws IOException {
    this.czArrowSession.flush();
  }

  @Override
  public MetricBase getMetrics() {
    return this.czArrowSession.getMetrics();
  }

  @Override
  public void close() throws IOException {
    if (this.czArrowSession != null) {
      this.czArrowSession.close();
    }
  }

  public String getServerToken() {
    return serverToken;
  }

  // Set the non-empty server token only if it is unset.
  // Throw CZException if the server token passed is not identical to what is cached.
  // For table that does not require server token, server token is always null or empty
  // so it will not fail the check.
  public void safeSetServerToken(String newServerToken) {
    safeSetServerToken(newServerToken, false);
  }

  public void safeSetServerToken(String newServerToken, boolean forceReset) {
    if (forceReset) {
      this.serverToken = newServerToken;
      LOG.info("setServerToken with new {} forceReset {}", newServerToken, forceReset);
      return;
    }
    newServerToken = (newServerToken != null && !newServerToken.isEmpty()) ? newServerToken : null;
    if (this.serverToken == null) {
      this.serverToken = newServerToken;
      LOG.info("setServerToken with new {} forceReset {}", newServerToken, forceReset);
    } else if (!this.serverToken.equals(newServerToken)) {
      throw new CZException(String.format("Table %s serverToken is not same, old: %s, new: %s",
          getTable().getTableName(), this.serverToken, newServerToken));
    }
  }
}