Search in sources :

Example 1 with BIGQUERY_FAILED_TO_EXECUTE_QUERY

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);
    }
}
Also used : BIGQUERY_FAILED_TO_EXECUTE_QUERY(io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_FAILED_TO_EXECUTE_QUERY) ConnectorSplitManager(io.trino.spi.connector.ConnectorSplitManager) Logger(io.airlift.log.Logger) IntStream.range(java.util.stream.IntStream.range) NodeManager(io.trino.spi.NodeManager) TableId(com.google.cloud.bigquery.TableId) BigQueryException(com.google.cloud.bigquery.BigQueryException) Duration(io.airlift.units.Duration) FixedSplitSource(io.trino.spi.connector.FixedSplitSource) Inject(javax.inject.Inject) NOT_SUPPORTED(io.trino.spi.StandardErrorCode.NOT_SUPPORTED) TableNotFoundException(io.trino.spi.connector.TableNotFoundException) ImmutableList(com.google.common.collect.ImmutableList) VIEW(com.google.cloud.bigquery.TableDefinition.Type.VIEW) ConnectorTableHandle(io.trino.spi.connector.ConnectorTableHandle) Objects.requireNonNull(java.util.Objects.requireNonNull) ColumnHandle(io.trino.spi.connector.ColumnHandle) TableResult(com.google.cloud.bigquery.TableResult) TABLE(com.google.cloud.bigquery.TableDefinition.Type.TABLE) ReadSession(com.google.cloud.bigquery.storage.v1.ReadSession) ImmutableList.toImmutableList(com.google.common.collect.ImmutableList.toImmutableList) TrinoException(io.trino.spi.TrinoException) ConnectorSplitSource(io.trino.spi.connector.ConnectorSplitSource) ConnectorSession(io.trino.spi.connector.ConnectorSession) TupleDomain(io.trino.spi.predicate.TupleDomain) SchemaTableName(io.trino.spi.connector.SchemaTableName) List(java.util.List) Collectors.toList(java.util.stream.Collectors.toList) DynamicFilter(io.trino.spi.connector.DynamicFilter) Optional(java.util.Optional) TableInfo(com.google.cloud.bigquery.TableInfo) ConnectorTransactionHandle(io.trino.spi.connector.ConnectorTransactionHandle) TableNotFoundException(io.trino.spi.connector.TableNotFoundException) TableResult(com.google.cloud.bigquery.TableResult) TrinoException(io.trino.spi.TrinoException) TableInfo(com.google.cloud.bigquery.TableInfo) BigQueryException(com.google.cloud.bigquery.BigQueryException) SchemaTableName(io.trino.spi.connector.SchemaTableName)

Aggregations

BigQueryException (com.google.cloud.bigquery.BigQueryException)1 TABLE (com.google.cloud.bigquery.TableDefinition.Type.TABLE)1 VIEW (com.google.cloud.bigquery.TableDefinition.Type.VIEW)1 TableId (com.google.cloud.bigquery.TableId)1 TableInfo (com.google.cloud.bigquery.TableInfo)1 TableResult (com.google.cloud.bigquery.TableResult)1 ReadSession (com.google.cloud.bigquery.storage.v1.ReadSession)1 ImmutableList (com.google.common.collect.ImmutableList)1 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)1 Logger (io.airlift.log.Logger)1 Duration (io.airlift.units.Duration)1 BIGQUERY_FAILED_TO_EXECUTE_QUERY (io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_FAILED_TO_EXECUTE_QUERY)1 NodeManager (io.trino.spi.NodeManager)1 NOT_SUPPORTED (io.trino.spi.StandardErrorCode.NOT_SUPPORTED)1 TrinoException (io.trino.spi.TrinoException)1 ColumnHandle (io.trino.spi.connector.ColumnHandle)1 ConnectorSession (io.trino.spi.connector.ConnectorSession)1 ConnectorSplitManager (io.trino.spi.connector.ConnectorSplitManager)1 ConnectorSplitSource (io.trino.spi.connector.ConnectorSplitSource)1 ConnectorTableHandle (io.trino.spi.connector.ConnectorTableHandle)1