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()
     {