com.clickzetta.client.jdbc.arrow.CZArrowCursor 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.arrow.accessor.ArrowAccessorFactory;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.FieldVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.calcite.avatica.ColumnMetaData;
import org.apache.calcite.avatica.util.AbstractCursor;
import org.apache.calcite.avatica.util.ArrayImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.Calendar;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * CZ Jdbc's Cursor class to visit arrow vectors.
 * 
 * Please refer to {@link org.apache.arrow.driver.jdbc.ArrowFlightJdbcCursor} for more details.
 */
public class CZArrowCursor extends AbstractCursor {

  private static final Logger LOGGER;

  static {
    LOGGER = LoggerFactory.getLogger(CZArrowCursor.class);
  }

  private final VectorSchemaRoot root;
  private final int rowCount;
  private int currentRow = -1;

  public CZArrowCursor(VectorSchemaRoot root) {
    this.root = root;
    rowCount = root.getRowCount();
  }

  @Override
  public List createAccessors(List columns,
                                        Calendar localCalendar,
                                        ArrayImpl.Factory factory) {
    final List fieldVectors = root.getFieldVectors();

    return IntStream.range(0, fieldVectors.size()).mapToObj(root::getVector)
      .map(this::createAccessor)
      .collect(Collectors.toCollection(() -> new ArrayList<>(fieldVectors.size())));
  }

  private Accessor createAccessor(FieldVector vector) {
    return ArrowAccessorFactory.createAccessor(vector, this::getCurrentRow,
      (boolean wasNull) -> {
        // AbstractCursor creates a boolean array of length 1 to hold the wasNull value
        this.wasNull[0] = wasNull;
      });
  }

  /**
   * ArrowFlightJdbcAccessors do not use {@link AbstractCursor.Getter}, as it would box primitive types and cause
   * performance issues. Each Accessor implementation works directly on Arrow Vectors.
   */
  @Override
  protected Getter createGetter(int column) {
    throw new UnsupportedOperationException("Not allowed.");
  }

  @Override
  public boolean next() {
    currentRow++;
    return currentRow < rowCount;
  }

  @Override
  public void close() {
    try {
      AutoCloseables.close(root);
    } catch (Exception e) {
      LOGGER.error(e.getMessage(), e);
    }
  }

  private int getCurrentRow() {
    return currentRow;
  }
}