use of io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_FAILED_TO_EXECUTE_QUERY in project trino by trinodb.
the class BigQuerySplitManager method createEmptyProjection.
private List<BigQuerySplit> createEmptyProjection(ConnectorSession session, TableId remoteTableId, int actualParallelism, Optional<String> filter) {
BigQueryClient client = bigQueryClientFactory.create(session);
log.debug("createEmptyProjection(tableId=%s, actualParallelism=%s, filter=[%s])", remoteTableId, actualParallelism, filter);
try {
long numberOfRows;
if (filter.isPresent()) {
// count the rows based on the filter
String sql = client.selectSql(remoteTableId, "COUNT(*)");
TableResult result = client.query(sql);
numberOfRows = result.iterateAll().iterator().next().get(0).getLongValue();
} else {
// no filters, so we can take the value from the table info when the object is TABLE
TableInfo tableInfo = client.getTable(remoteTableId).orElseThrow(() -> new TableNotFoundException(new SchemaTableName(remoteTableId.getDataset(), remoteTableId.getTable())));
if (tableInfo.getDefinition().getType() == TABLE) {
numberOfRows = tableInfo.getNumRows().longValue();
} else if (tableInfo.getDefinition().getType() == VIEW) {
String sql = client.selectSql(remoteTableId, "COUNT(*)");
TableResult result = client.query(sql);
numberOfRows = result.iterateAll().iterator().next().get(0).getLongValue();
} else {
throw new TrinoException(NOT_SUPPORTED, "Unsupported table type: " + tableInfo.getDefinition().getType());
}
}
long rowsPerSplit = numberOfRows / actualParallelism;
// need to be added to one fo the split due to integer division
long remainingRows = numberOfRows - (rowsPerSplit * actualParallelism);
List<BigQuerySplit> splits = range(0, actualParallelism).mapToObj(ignored -> BigQuerySplit.emptyProjection(rowsPerSplit)).collect(toList());
splits.set(0, BigQuerySplit.emptyProjection(rowsPerSplit + remainingRows));
return splits;
} catch (BigQueryException e) {
throw new TrinoException(BIGQUERY_FAILED_TO_EXECUTE_QUERY, "Failed to compute empty projection", e);
}
}
Aggregations