com.clickzetta.client.jdbc.arrow.CZArrowResultSet 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.client.jdbc.arrow;

import com.clickzetta.client.jdbc.core.CZColumnMetaData;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.calcite.avatica.AvaticaResultSet;
import org.apache.calcite.avatica.AvaticaResultSetMetaData;
import org.apache.calcite.avatica.AvaticaStatement;
import org.apache.calcite.avatica.ColumnMetaData;
import org.apache.calcite.avatica.Meta;
import org.apache.calcite.avatica.QueryState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.util.*;

import static java.util.Objects.isNull;

/**
 * ResultSet implementation for a set of Arrow Vectors.
 */
public class CZArrowResultSet extends AvaticaResultSet {

  private static final Logger logger = LoggerFactory.getLogger(CZArrowResultSet.class);
  private int currentRow = 0;
  private final int maxRows;
  private List columnMetaDatas;
  private CloseableVectorSchemaRootIterator vectorSchemaRootIterator;
  private VectorSchemaRoot currentVectorSchemaRoot;
  private List locations = new ArrayList<>();

  CZArrowResultSet(final AvaticaStatement statement,
                   final QueryState state,
                   final Meta.Signature signature,
                   final ResultSetMetaData resultSetMetaData,
                   final TimeZone timeZone,
                   final Meta.Frame firstFrame,
                   int maxRows) throws SQLException {
    super(statement, state, signature, resultSetMetaData, timeZone, firstFrame);
    this.maxRows = maxRows;
  }

  public static CZArrowResultSet fromIterator(final List czColumnMetaDatas,
                                              final CloseableVectorSchemaRootIterator iterator) throws SQLException {
    return fromIterator(czColumnMetaDatas, iterator, 10000);
  }

  public static CZArrowResultSet fromIterator(final List czColumnMetaDatas,
                                              final CloseableVectorSchemaRootIterator iterator,
                                              int maxRows) throws SQLException {
    List columnMetaDatas = ConvertUtils.convertToColumnMetaDataList(czColumnMetaDatas);

    // Similar to how org.apache.calcite.avatica.util.ArrayFactoryImpl does
    final TimeZone timeZone = TimeZone.getDefault();
    final QueryState state = new QueryState();
    final Meta.Signature signature = new Meta.Signature(
      columnMetaDatas,
      null,
      Collections.emptyList(),
      Collections.emptyMap(),
      null, // unnecessary, as SQL requests use ArrowCursor
      Meta.StatementType.SELECT
    );

    final AvaticaResultSetMetaData resultSetMetaData =
      new AvaticaResultSetMetaData(null, null, signature);
    final CZArrowResultSet resultSet = new CZArrowResultSet(
      null, state, signature, resultSetMetaData, timeZone, null, maxRows);
    resultSet.execute(columnMetaDatas, iterator);
    logger.info("Creating CZArrowResultSet from iterator, this: " + resultSet);
    return resultSet;
  }

  @Override
  protected AvaticaResultSet execute() throws SQLException {
    throw new RuntimeException("Can only execute with execute(Iterator)");
  }

  private void execute(final VectorSchemaRoot vectorSchemaRoot) {
    execute2(new CZArrowCursor(vectorSchemaRoot), columnMetaDatas);
  }

  void execute(final List columnMetaDatas, final CloseableVectorSchemaRootIterator iterator) {
    this.columnMetaDatas = columnMetaDatas;
    this.vectorSchemaRootIterator = iterator;
    loadNextVectorSchemaRoot();
    executeForCurrentVectorSchemaRoot();
  }

  private void loadNextVectorSchemaRoot() {
    if (vectorSchemaRootIterator == null) {
      return;
    }
    if (vectorSchemaRootIterator.hasNext()) {
      currentVectorSchemaRoot = vectorSchemaRootIterator.next();
    } else {
      try {
        closeVectors();
      } catch (final Exception e) {
        throw new RuntimeException(e);
      }
    }
  }

  private void executeForCurrentVectorSchemaRoot() {
    if (currentVectorSchemaRoot != null) {
      execute(currentVectorSchemaRoot);
    }
  }

  // Close the result set internally to make sure no further rows are read.
  private void closeInternal() {
    try {
      logger.info("Start closing result set {} internal.", this);
      if (statement != null && statement.isCloseOnCompletion()) {
        statement.close();
      }
      closeVectors();
      // If no current row, zero it to conform JDBC specs.
      currentRow = 0;
    } catch (final Exception e) {
      throw new RuntimeException(e);
    }
  }

  private void closeVectors() throws Exception {
    AutoCloseables.close(currentVectorSchemaRoot, vectorSchemaRootIterator);
    currentVectorSchemaRoot = null;
    vectorSchemaRootIterator = null;
  }

  @Override
  public boolean next() throws SQLException {
    if (vectorSchemaRootIterator == null) {
      return false;
    }

    while (true) {
      final boolean hasNext = super.next();
      if (maxRows != 0 && currentRow >= maxRows) {
        closeInternal();
        return false;
      }

      if (hasNext) {
        currentRow++;
        return true;
      }

      loadNextVectorSchemaRoot();
      if (currentVectorSchemaRoot != null) {
        executeForCurrentVectorSchemaRoot();
        continue;
      }

      closeInternal();
      return false;
    }
  }

  @Override
  protected void cancel() {
    signature.columns.clear();
    super.cancel();
    try {
      closeVectors();
      currentRow = 0;
    } catch (final Exception e) {
      throw new RuntimeException(e);
    }
  }

  @Override
  public void close() {
    logger.info("Start closing CZArrowResultSet, this: " + this);
    final Set exceptions = new HashSet<>();
    try {
      if (isClosed()) {
        return;
      }
    } catch (final SQLException e) {
      exceptions.add(e);
    }
    try {
      closeVectors();
      currentRow = 0;
    } catch (final Exception e) {
      exceptions.add(e);
    }
    if (!isNull(statement)) {
      try {
        super.close();
      } catch (final Exception e) {
        exceptions.add(e);
      }
    }
    exceptions.parallelStream().forEach(e -> logger.error(e.getMessage(), e));
    exceptions.stream().findAny().ifPresent(e -> {
      throw new RuntimeException(e);
    });
  }

  // We need to override getObject as the AvaticaResultSet implementation
  // has unexpected results.
  @Override
  public Object getObject(int columnIndex) throws SQLException {
    checkOpen();
    try {
      return accessorList.get(columnIndex - 1).getObject();
    } catch (IndexOutOfBoundsException e) {
      throw new SQLException("invalid column ordinal: " + columnIndex);
    }
  }

  @Override
  public String getString(int columnIndex) throws SQLException {
    String str = super.getString(columnIndex);
    if (str != null) {
      return str;
    }
    return "null";
  }

  @Override
  public String getString(String columnLabel) throws SQLException {
    String str = super.getString(columnLabel);
    if (str != null) {
      return str;
    }
    return "null";
  }

  @Override
  public int getRow() throws SQLException {
    checkOpen();
    return currentRow;
  }

  @Override
  public int getType() throws SQLException {
    return TYPE_FORWARD_ONLY;
  }

  @Override
  public boolean isFirst() throws SQLException {
    checkOpen();
    return currentRow == 1;
  }

  public List getLocations() {
    return locations;
  }

  public void setLocations(List locations) {
    this.locations = locations;
  }

  public void addLocations(String location) {
    this.locations.add(location);
  }
}