Fix wrong result due to column position mismatch in BigQuery

trinodb · Jun 6, 2021 · 1fc5895 · 1fc5895
1 parent 9ad26b4
commit 1fc5895
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryResultPageSource.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryResultPageSource.java
@@ -76,6 +76,7 @@ public class BigQueryResultPageSource
 
  private final BigQueryStorageClient bigQueryStorageClient;
  private final BigQuerySplit split;
+ private final List<String> columnNames;
  private final List<Type> columnTypes;
  private final AtomicLong readBytes;
  private final PageBuilder pageBuilder;
@@ -90,7 +91,11 @@ public BigQueryResultPageSource(
  this.bigQueryStorageClient = requireNonNull(bigQueryStorageClientFactory, "bigQueryStorageClientFactory is null").createBigQueryStorageClient();
  this.split = requireNonNull(split, "split is null");
  this.readBytes = new AtomicLong();
- this.columnTypes = requireNonNull(columns, "columns is null").stream()
+ requireNonNull(columns, "columns is null");
+ this.columnNames = columns.stream()
+ .map(BigQueryColumnHandle::getName)
+ .collect(toImmutableList());
+ this.columnTypes = columns.stream()
  .map(BigQueryColumnHandle::getTrinoType)
  .collect(toImmutableList());
  this.pageBuilder = new PageBuilder(columnTypes);
@@ -131,7 +136,7 @@ public Page getNextPage()
  pageBuilder.declarePosition();
  for (int column = 0; column < columnTypes.size(); column++) {
  BlockBuilder output = pageBuilder.getBlockBuilder(column);
- appendTo(columnTypes.get(column), record.get(column), output);
+ appendTo(columnTypes.get(column), record.get(columnNames.get(column)), output);
  }
  }
 

diff --git a/...ino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java b/...ino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java
@@ -317,6 +317,26 @@ public void testRepeatCountAggregationView()
  onBigQuery("DROP VIEW " + viewName);
  }
 
+ /**
+ * https://github.com/trinodb/trino/issues/8183
+ */
+ @Test
+ public void testColumnPositionMismatch()
+ {
+ String tableName = "test.test_column_position_mismatch";
+
+ assertUpdate("DROP TABLE IF EXISTS " + tableName);
+ assertUpdate(format("CREATE TABLE %s (c_varchar VARCHAR, c_int INT, c_date DATE)", tableName));
+ onBigQuery(format("INSERT INTO %s VALUES ('a', 1, '2021-01-01')", tableName));
+
+ // Adding a CAST makes BigQuery return columns in a different order
+ assertQuery(
+ "SELECT c_varchar, CAST(c_int AS SMALLINT), c_date FROM " + tableName,
+ "VALUES ('a', 1, '2021-01-01')");
+
+ assertUpdate("DROP TABLE " + tableName);
+ }
+
  @Test
  public void testViewDefinitionSystemTable()
  {