Examples with VIEW com.google.cloud.bigquery.TableDefinition.Type.VIEW used on opensource projects

Example 1 with VIEW

use of com.google.cloud.bigquery.TableDefinition.Type.VIEW 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);
    }
}