com.clickzetta.client.jdbc.arrow.CZInMemoryArrowResult 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.memory.RootAllocator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Iterator;
import java.util.List;

public class CZInMemoryArrowResult extends CZArrowResult {

  private static final Logger logger = LoggerFactory.getLogger(CZInMemoryArrowResult.class);
  private final Iterator inputIterator;

  public CZInMemoryArrowResult(List inMemoryArrowFiles) {
    this(inMemoryArrowFiles, new RootAllocator(Long.MAX_VALUE));
  }

  public CZInMemoryArrowResult(List inMemoryArrowFiles, BufferAllocator allocator) {
    super(allocator);
    this.inputIterator = inMemoryArrowFiles.iterator();
    this.allocator = allocator;
  }

  @Override
  protected boolean hasNextFile() {
    return inputIterator.hasNext();
  }

  @Override
  protected ArrowVectorIterator readNextFile() {
    try {
      logger.debug("Read next in-memory arrow ipc file.");
      // The returned ArrowVectorIterator will be closed by the super class CZArrowResult
      return new ArrowVectorIterator(inputIterator.next(), allocator);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  @Override
  public void close() throws Exception {
    super.close();
  }
}