com.clickzetta.client.jdbc.arrow.CZArrowFileResult 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.CZStorageClient;
import cz.proto.coordinator.CoordinatorServiceOuterClass;
import org.apache.arrow.memory.BufferAllocator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.io.InputStream;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

public class CZArrowFileResult extends CZArrowResult {

  private static final Logger logger = LoggerFactory.getLogger(CZArrowFileResult.class);

  private Iterator fileIterator;
  private CZStorageClient storageClient;

  public CZArrowFileResult(CoordinatorServiceOuterClass.JobResultLocation jobResultLocation,
                           boolean useInternalEndpoint,
                           BufferAllocator allocator,
                           boolean useObjectStoreHttps) {
    super(allocator);

    List files = new ArrayList<>(jobResultLocation.getLocationCount());
    for (int i = 0; i < jobResultLocation.getLocationCount(); i++) {
      files.add(jobResultLocation.getLocation(i));
    }
    if (files.isEmpty()) {
      logger.warn("no file in JobResultLocation.");
      return;
    }
    fileIterator = files.iterator();
    storageClient = CZStorageClient.create(jobResultLocation, useInternalEndpoint, useObjectStoreHttps);
    logger.info("Init storage client done, storage client: " + storageClient);
  }

  // Testing purpose
  public CZArrowFileResult(List files, CZStorageClient storageClient, BufferAllocator allocator) {
    super(allocator);
    fileIterator = files.iterator();
    this.storageClient = storageClient;
  }

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

  @Override
  protected ArrowVectorIterator readNextFile() {
    String nextFile = fileIterator.next();
    logger.info("ReadFileLocation: " + nextFile);
    URI uri = null;
    try {
      uri = new URI(nextFile);
    } catch (URISyntaxException e) {
      throw new RuntimeException(e);
    }
    String bucketName = uri.getHost();
    String file = uri.getPath().substring(1);
    logger.info("ObjectStorageBucket: " + bucketName);
    logger.info("ObjectStorageFile: " + file);
    try {
      // The returned ArrowVectorIterator will be closed by the super class CZArrowResult
      InputStream inputStream = storageClient.getObjectInputStream(bucketName, file);
      logger.info("Get object input stream from storage client with file: " + file + ", inputStream: " + inputStream);
      return new ArrowVectorIterator(inputStream, allocator, file);
    } catch (IOException e) {
      throw new RuntimeException(e);
    }
  }

  private void closeStorageClient() {
    if (storageClient != null) {
      String storageClientStr = storageClient.toString();
      storageClient.shutDown();
      storageClient = null;
      logger.info("Storage client {} is closed.", storageClientStr);
    }
  }

  @Override
  public void close() throws Exception {
    logger.info("Start closing arrow file result set: {}", this);
    Exception exception = null;

    try {
      super.close();
    } catch (Exception e) {
      exception = e;
    }

    try {
      closeStorageClient();
    } catch (Exception e) {
      if (exception != null) {
        exception.addSuppressed(e);
      } else {
        exception = e;
      }
    }

    if (exception != null) {
      throw exception;
    }
  }

}