com.clickzetta.client.jdbc.arrow.ArrowVectorIterator 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 org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.util.AutoCloseables;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowStreamReader;
import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;

public class ArrowVectorIterator implements CloseableVectorSchemaRootIterator {
  private static final Logger logger = LoggerFactory.getLogger(ArrowVectorIterator.class);

  private final ArrowStreamReader reader;
  private final InputStream inputStream;
  private VectorSchemaRoot currentVectorSchemaRoot;
  private String file;

  // Read a set of VectorSchemaRoot from input stream of an arrow feather stream file.
  public ArrowVectorIterator(InputStream inputStream, BufferAllocator allocator, String file) throws IOException {
    this.inputStream = inputStream;
    this.reader = new ArrowStreamReader(inputStream, allocator);
    this.file = file;
  }

  // Read a set of VectorSchemaRoot from an in-memory byte array of an arrow feather stream file.
  public ArrowVectorIterator(byte[] bytes, BufferAllocator allocator) throws IOException {
    this.inputStream = null;
    this.reader = new ArrowStreamReader(new ByteArrayReadableSeekableByteChannel(bytes), allocator);
  }

  @Override
  public boolean hasNext() {
    try {
      AutoCloseables.close(currentVectorSchemaRoot);
      currentVectorSchemaRoot = null;

      boolean hasNextBatch = reader.loadNextBatch();
      if (hasNextBatch) {
        currentVectorSchemaRoot = reader.getVectorSchemaRoot();
      }
      return hasNextBatch;
    } catch (Exception e) {
      throw new RuntimeException(e);
    }
  }

  @Override
  public VectorSchemaRoot next() {
    return currentVectorSchemaRoot;
  }

  @Override
  public void close() throws Exception {
    logger.info("Start closing input stream and reader for file: " + file + ", inputStream: " + inputStream);
    AutoCloseables.close(currentVectorSchemaRoot, reader, inputStream);
  }
}