Skip to content

Commit

Permalink
Skip unsupported timestamp type in BigQuery
Browse files Browse the repository at this point in the history
  • Loading branch information
@ebyhr
ebyhr committed Nov 7, 2024
1 parent 2f607e1 commit 8be3ee7
Show file tree
Hide file tree
Showing 5 changed files with 39 additions and 13 deletions.
8 changes: 4 additions & 4 deletions plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -565,10 +565,10 @@ public List<BigQueryColumnHandle> getColumns(BigQueryTableHandle tableHandle)

TableInfo tableInfo = getTable(tableHandle.asPlainTable().getRemoteTableName().toTableId())
.orElseThrow(() -> new TableNotFoundException(tableHandle.asPlainTable().getSchemaTableName()));
return buildColumnHandles(tableInfo);
return buildColumnHandles(tableInfo, tableHandle.relationHandle().isUseStorageApi());
}

public List<BigQueryColumnHandle> buildColumnHandles(TableInfo tableInfo)
public List<BigQueryColumnHandle> buildColumnHandles(TableInfo tableInfo, boolean useStorageApi)
{
Schema schema = tableInfo.getDefinition().getSchema();
if (schema == null) {
Expand All @@ -577,8 +577,8 @@ public List<BigQueryColumnHandle> buildColumnHandles(TableInfo tableInfo)
}
return schema.getFields()
.stream()
.filter(typeManager::isSupportedType)
.map(typeManager::toColumnHandle)
.filter(field -> typeManager.isSupportedType(field, useStorageApi))
.map(field -> typeManager.toColumnHandle(field, useStorageApi))
.collect(toImmutableList());
}

Expand Down
7 changes: 4 additions & 3 deletions plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -340,8 +340,9 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable
return null;
}

boolean useStorageApi = useStorageApi(session, schemaTableName.getTableName(), tableInfo.get().getDefinition().getType());
ImmutableList.Builder<BigQueryColumnHandle> columns = ImmutableList.builder();
columns.addAll(client.buildColumnHandles(tableInfo.get()));
columns.addAll(client.buildColumnHandles(tableInfo.get(), useStorageApi));
Optional<BigQueryPartitionType> partitionType = getPartitionType(tableInfo.get().getDefinition());
if (partitionType.isPresent() && partitionType.get() == INGESTION) {
columns.add(PARTITION_DATE.getColumnHandle());
Expand All @@ -353,7 +354,7 @@ public ConnectorTableHandle getTableHandle(ConnectorSession session, SchemaTable
tableInfo.get().getDefinition().getType().toString(),
partitionType,
Optional.ofNullable(tableInfo.get().getDescription()),
useStorageApi(session, schemaTableName.getTableName(), tableInfo.get().getDefinition().getType())),
useStorageApi),
TupleDomain.all(),
Optional.empty())
.withProjectedColumns(columns.build());
Expand Down Expand Up @@ -469,7 +470,7 @@ public Map<SchemaTableName, List<ColumnMetadata>> listTableColumns(ConnectorSess
return tableInfos.stream()
.collect(toImmutableMap(
table -> new SchemaTableName(table.getTableId().getDataset(), table.getTableId().getTable()),
table -> client.buildColumnHandles(table).stream()
table -> client.buildColumnHandles(table, useStorageApi(session, table.getTableId().getTable(), table.getDefinition().getType())).stream()
.map(BigQueryColumnHandle::getColumnMetadata)
.collect(toImmutableList())));
}
Expand Down
12 changes: 8 additions & 4 deletions plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryTypeManager.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -380,14 +380,14 @@ private Optional<ColumnMapping> convertToTrinoType(Field field)
}
}

public BigQueryColumnHandle toColumnHandle(Field field)
public BigQueryColumnHandle toColumnHandle(Field field, boolean useStorageApi)
{
FieldList subFields = field.getSubFields();
List<BigQueryColumnHandle> subColumns = subFields == null ?
Collections.emptyList() :
subFields.stream()
.filter(this::isSupportedType)
.map(this::toColumnHandle)
.filter(column -> isSupportedType(column, useStorageApi))
.map(column -> toColumnHandle(column, useStorageApi))
.collect(Collectors.toList());
ColumnMapping columnMapping = toTrinoType(field).orElseThrow(() -> new IllegalArgumentException("Unsupported type: " + field));
return new BigQueryColumnHandle(
Expand All @@ -402,7 +402,7 @@ public BigQueryColumnHandle toColumnHandle(Field field)
false);
}

public boolean isSupportedType(Field field)
public boolean isSupportedType(Field field, boolean useStorageApi)
{
LegacySQLTypeName type = field.getType();
if (type == LegacySQLTypeName.BIGNUMERIC) {
Expand All @@ -414,6 +414,10 @@ public boolean isSupportedType(Field field)
return false;
}
}
if (!useStorageApi && type == LegacySQLTypeName.TIMESTAMP) {
// TODO https://github.com/trinodb/trino/issues/12346 BigQueryQueryPageSource does not support TIMESTAMP type
return false;
}

return toTrinoType(field).isPresent();
}
Expand Down
4 changes: 2 additions & 2 deletions plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/ptf/Query.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,11 +122,11 @@ public TableFunctionAnalysis analyze(

ImmutableList.Builder<BigQueryColumnHandle> columnsBuilder = ImmutableList.builderWithExpectedSize(schema.getFields().size());
for (com.google.cloud.bigquery.Field field : schema.getFields()) {
if (!typeManager.isSupportedType(field)) {
if (!typeManager.isSupportedType(field, useStorageApi)) {
// TODO: Skip unsupported type instead of throwing an exception
throw new TrinoException(NOT_SUPPORTED, "Unsupported type: " + field.getType());
}
columnsBuilder.add(typeManager.toColumnHandle(field));
columnsBuilder.add(typeManager.toColumnHandle(field, useStorageApi));
}

Descriptor returnedType = new Descriptor(columnsBuilder.build().stream()
Expand Down
21 changes: 21 additions & 0 deletions plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/BaseBigQueryConnectorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import io.trino.testing.QueryRunner.MaterializedResultWithPlan;
import io.trino.testing.TestingConnectorBehavior;
import io.trino.testing.sql.TestTable;
import io.trino.testing.sql.TestView;
import org.intellij.lang.annotations.Language;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Disabled;
Expand Down Expand Up @@ -644,6 +645,26 @@ public void testSkipUnsupportedType()
}
}

@Test
public void testSkipUnsupportedTimestampType()
{
Session skipViewMaterialization = Session.builder(getSession())
.setCatalogSessionProperty("bigquery", "skip_view_materialization", "true")
.build();

try (TestView view = new TestView(
bigQuerySqlExecutor,
"test.test_skip_unsupported_type",
"SELECT 1 a, TIMESTAMP '1970-01-01 00:00:00 UTC' unsupported, 2 b")) {
assertQuery(skipViewMaterialization, "SELECT * FROM " + view.getName(), "VALUES (1, 2)");
assertThat((String) computeActual(skipViewMaterialization, "SHOW CREATE TABLE " + view.getName()).getOnlyValue())
.isEqualTo("CREATE TABLE bigquery." + view.getName() + " (\n" +
" a bigint,\n" +
" b bigint\n" +
")");
}
}

@Test
@Override
public void testDateYearOfEraPredicate()
Expand Down

0 comments on commit 8be3ee7

Please sign in to comment.