com.clickzetta.client.jdbc.core.CZLakehouseResponseHelper 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
Show all versions of clickzetta-java Show documentation
The java SDK for clickzetta's Lakehouse
package com.clickzetta.client.jdbc.core;;
import com.clickzetta.client.jdbc.arrow.CZArrowFileResult;
import com.clickzetta.client.jdbc.arrow.CZArrowResult;
import com.clickzetta.client.jdbc.arrow.CZArrowResultSet;
import com.clickzetta.client.jdbc.arrow.CZInMemoryArrowResult;
import com.clickzetta.client.jdbc.arrow.util.CZTimestamp;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.util.JsonFormat;
import cz.proto.coordinator.CoordinatorServiceOuterClass;
import org.apache.arrow.memory.RootAllocator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.File;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
public class CZLakehouseResponseHelper {
static Logger logger = LoggerFactory.getLogger(CZLakehouseResponseHelper.class);
static protected List parseColumnMetaData(CoordinatorServiceOuterClass.JobResultSetMetadata metadata)
throws SQLException {
List columnMetaDatas = new ArrayList<>(metadata.getFieldsCount());
if (metadata.getFormat() == CoordinatorServiceOuterClass.ResultFormat.TEXT
|| metadata.getFormat() == CoordinatorServiceOuterClass.ResultFormat.ARROW
|| metadata.getFormat() == CoordinatorServiceOuterClass.ResultFormat.CSV) {
for (int i = 0; i < metadata.getFieldsCount(); i++) {
cz.proto.FieldSchema fieldSchema = metadata.getFields(i);
columnMetaDatas.add(new CZColumnMetaData(fieldSchema));
}
} else {
throw new SQLException("UnsupportedResultFormat: " + metadata.getFormat());
}
return columnMetaDatas;
}
public static CoordinatorServiceOuterClass.SubmitJobResponse convertToPb(String response) throws Exception {
CoordinatorServiceOuterClass.SubmitJobResponse.Builder respBuilder =
CoordinatorServiceOuterClass.SubmitJobResponse.newBuilder();
JsonFormat.parser().ignoringUnknownFields().merge(response, respBuilder);
CoordinatorServiceOuterClass.SubmitJobResponse resp = respBuilder.build();
return resp;
}
public static ResultSet parseSubmitJobResponse(CoordinatorServiceOuterClass.SubmitJobResponse submitJobResponse,
Map configs) throws Exception {
return parseSubmitJobResponse(submitJobResponse, 10000, configs);
}
public static ResultSet parseSubmitJobResponse(CoordinatorServiceOuterClass.SubmitJobResponse submitJobResponse,
int maxRowSize, Map configs) throws Exception {
return parseSubmitJobResponse(submitJobResponse, maxRowSize, true, configs, true);
}
public static ResultSet parseSubmitJobResponse(CoordinatorServiceOuterClass.SubmitJobResponse submitJobResponse,
int maxRowSize, boolean useInternalOssEndpoint,
Map configs, boolean useObjectStoreHttps) throws Exception {
return parseJobResult(submitJobResponse.getStatus(), submitJobResponse.getResultSet(), maxRowSize,
useInternalOssEndpoint, configs, useObjectStoreHttps);
}
public static ResultSet parseJobResult(CoordinatorServiceOuterClass.JobStatus jobStatus,
CoordinatorServiceOuterClass.JobResultSet resultSet,
int maxRowSize, boolean useInternalOssEndpoint,
Map configs,
boolean useObjectStoreHttps) throws Exception {
long startMs = System.currentTimeMillis();
String state = jobStatus.getState();
if (state.isEmpty() || state.equals("FAILED") || state.equals("CANCELLING") || state.equals("CANCELLED")) {
throw new Exception(jobStatus.getMessage());
}
CoordinatorServiceOuterClass.JobResultSetMetadata metadata = resultSet.getMetadata();
List columnMetaDatas = new ArrayList<>();
if (resultSet.hasMetadata()) {
columnMetaDatas = parseColumnMetaData(metadata);
} else {
return null;
}
ResultSet result;
if (metadata.getFormat() == CoordinatorServiceOuterClass.ResultFormat.TEXT) {
result = parseTextResult(resultSet, columnMetaDatas, maxRowSize, useInternalOssEndpoint, useObjectStoreHttps);
} else if (metadata.getFormat() == CoordinatorServiceOuterClass.ResultFormat.ARROW) {
result = parseArrowResult(resultSet, columnMetaDatas, maxRowSize,
useInternalOssEndpoint, configs, useObjectStoreHttps);
} else if (metadata.getFormat() == CoordinatorServiceOuterClass.ResultFormat.CSV) {
result = parseCsvResult(resultSet, columnMetaDatas, maxRowSize, useInternalOssEndpoint, useObjectStoreHttps);
} else {
throw new Exception("UnsupportedResultFormat: " + metadata.getFormat());
}
logger.info("ParseResultCost(ms): {}", System.currentTimeMillis() - startMs);
return result;
}
private static ResultSet parseTextResult(CoordinatorServiceOuterClass.JobResultSet resultSet,
List columnMetaDatas,
int maxRowSize,
boolean useInternalOssEndpoint,
boolean useObjectStoreHttps) throws Exception {
CZResult czResult;
if (resultSet.hasData()) {
List rawDatas = new ArrayList<>();
for (int i = 0; i < resultSet.getData().getDataCount(); i++) {
String decoded = resultSet.getData().getData(i).toStringUtf8();
rawDatas.add(decoded);
}
czResult = new CZInMemoryTextResult(rawDatas, columnMetaDatas, maxRowSize);
} else if (resultSet.hasLocation()) {
logger.info("location count: {}", resultSet.getLocation().getLocationCount());
if (resultSet.getLocation().getLocationCount() > 0) {
czResult = new CZTextFileResult(resultSet.getLocation(), columnMetaDatas, useInternalOssEndpoint,
maxRowSize, useObjectStoreHttps);
} else {
czResult = new CZInMemoryTextResult(new ArrayList<>(), columnMetaDatas, maxRowSize);
}
} else {
czResult = new CZInMemoryTextResult(new ArrayList<>(), columnMetaDatas, maxRowSize);
}
return new CZTextResultSet(columnMetaDatas, czResult);
}
private static ResultSet parseArrowResult(CoordinatorServiceOuterClass.JobResultSet resultSet,
List columnMetaDatas,
int maxRowSize, boolean useInternalOssEndpoint,
Map configs,
boolean useObjectStoreHttps) throws SQLException {
CZTimestamp.resetSessionTimezone(configs);
RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
CZArrowResult czArrowResult;
if (resultSet.hasData()) {
List blocks = new ArrayList<>();
resultSet.getData().getDataList().stream().map(block -> block.toByteArray()).forEach(blocks::add);
czArrowResult = new CZInMemoryArrowResult(blocks, allocator);
} else if (resultSet.hasLocation()) {
logger.info("location count: {}", resultSet.getLocation().getLocationCount());
if (resultSet.getLocation().getLocationCount() > 0) {
czArrowResult = new CZArrowFileResult(resultSet.getLocation(),
useInternalOssEndpoint, allocator, useObjectStoreHttps);
} else {
czArrowResult = new CZInMemoryArrowResult(new ArrayList<>(), allocator);
}
} else {
czArrowResult = new CZInMemoryArrowResult(new ArrayList<>(), allocator);
}
return CZArrowResultSet.fromIterator(columnMetaDatas, czArrowResult, maxRowSize);
}
private static ResultSet parseCsvResult(CoordinatorServiceOuterClass.JobResultSet resultSet,
List columnMetaDatas,
int maxRowSize,
boolean useInternalOssEndpoint,
boolean useObjectStoreHttps) throws Exception {
CZResult czResult;
if (resultSet.hasData()) {
List rawDatas = new ArrayList<>();
for (int i = 0; i < resultSet.getData().getDataCount(); i++) {
String decoded = resultSet.getData().getData(i).toStringUtf8();
rawDatas.add(decoded);
}
czResult = new CZInMemoryTextResult(rawDatas, columnMetaDatas, maxRowSize, CZLazySimpleSerDe.FileFormat.CSV);
} else if (resultSet.hasLocation()) {
logger.info("location count: {}", resultSet.getLocation().getLocationCount());
if (resultSet.getLocation().getLocationCount() > 0) {
czResult = new CZTextFileResult(resultSet.getLocation(),
columnMetaDatas,
useInternalOssEndpoint,
maxRowSize,
CZLazySimpleSerDe.FileFormat.CSV,
useObjectStoreHttps);
} else {
czResult = new CZInMemoryTextResult(new ArrayList<>(),
columnMetaDatas,
maxRowSize,
CZLazySimpleSerDe.FileFormat.CSV);
}
} else {
czResult = new CZInMemoryTextResult(new ArrayList<>(),
columnMetaDatas,
maxRowSize,
CZLazySimpleSerDe.FileFormat.CSV);
}
CZTextResultSet result = new CZTextResultSet(columnMetaDatas, czResult);
for (int i = 0; i < resultSet.getLocation().getPresignedUrlsCount(); i++) {
result.addLocation(resultSet.getLocation().getPresignedUrls(i));
}
return result;
}
public static ResultSet parseSubmitJobResponse(String response) throws Exception {
return parseSubmitJobResponse(response, 10000, true, true);
}
public static ResultSet parseSubmitJobResponse(String response, int maxRowSize,
boolean useInternalOssEndpoint,
boolean useObjectStoreHttps) throws Exception {
CoordinatorServiceOuterClass.SubmitJobResponse resp = convertToPb(response);
if (resp.getStatus().getState().isEmpty()) {
throw new Exception("Not success: " + response);
}
return parseJobResult(resp.getStatus(), resp.getResultSet(), maxRowSize,
useInternalOssEndpoint, null, useObjectStoreHttps);
}
public static ResultSet parseGetJobResultResponse(String response, boolean useInternalOssEndpoint,
boolean useObjectStoreHttps) throws Exception {
CoordinatorServiceOuterClass.GetJobResultResponse.Builder respBuilder =
CoordinatorServiceOuterClass.GetJobResultResponse.newBuilder();
JsonFormat.parser().ignoringUnknownFields().merge(response, respBuilder);
CoordinatorServiceOuterClass.GetJobResultResponse resp = respBuilder.build();
if (resp.getStatus().getState().isEmpty()) {
throw new Exception("Not success: " + response);
}
return parseJobResult(resp.getStatus(), resp.getResultSet(), 0,
useInternalOssEndpoint, null, useObjectStoreHttps);
}
public static ResultSet genVolumeResultSet(String srcFile, String destFile, CZStatement.CommandType commandType)
throws SQLException {
File file;
if (commandType.equals(CZStatement.CommandType.PUT_VOLUME)) {
file = new File(srcFile);
srcFile = file.getName();
} else if (commandType.equals(CZStatement.CommandType.GET_VOLUME)) {
file = new File(destFile);
destFile = file.getName();
} else {
throw new SQLException("Unsupported commandType: " + commandType);
}
List columnMetaDatas = new ArrayList<>();
columnMetaDatas.add(new CZColumnMetaData("source", "STRING", java.sql.Types.VARCHAR));
columnMetaDatas.add(new CZColumnMetaData("target", "STRING", java.sql.Types.VARCHAR));
columnMetaDatas.add(new CZColumnMetaData("source_size", "INT64", java.sql.Types.BIGINT));
columnMetaDatas.add(new CZColumnMetaData("target_size", "INT64", java.sql.Types.BIGINT));
// TODO: add source_compression, source_auto_compress
// columnMetaDatas.add(new CZColumnMetaData("source_compression", "STRING", java.sql.Types.VARCHAR));
// columnMetaDatas.add(new CZColumnMetaData("source_auto_compress", "BOOLEAN", java.sql.Types.BOOLEAN));
columnMetaDatas.add(new CZColumnMetaData("status", "STRING", java.sql.Types.VARCHAR));
List rawDatas = new ArrayList<>();
rawDatas.add(String.format("%s,%s,%d,%d,%s",
srcFile, destFile, file.length(), file.length(), "SUCCEED"));
CZResult czResult = new CZInMemoryTextResult(rawDatas, columnMetaDatas, 1);
return new CZTextResultSet(columnMetaDatas, czResult);
}
}
© 2015 - 2024 Weber Informatics LLC | Privacy Policy