Fix incorrect result when aggregating count BigQuery view

ebyhr · ebyhr · commit 3b1c23ee98f4 · 2020-12-21T10:40:21.000+09:00
diff --git a/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQuerySplitManager.java b/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQuerySplitManager.java
@@ -15,6 +15,7 @@
 
 import com.google.cloud.bigquery.BigQueryException;
 import com.google.cloud.bigquery.TableId;
+import com.google.cloud.bigquery.TableInfo;
 import com.google.cloud.bigquery.TableResult;
 import com.google.cloud.bigquery.storage.v1beta1.Storage.ReadSession;
 import com.google.common.collect.ImmutableList;
@@ -37,8 +38,11 @@
 import java.util.Optional;
 import java.util.OptionalInt;
 
+import static com.google.cloud.bigquery.TableDefinition.Type.TABLE;
+import static com.google.cloud.bigquery.TableDefinition.Type.VIEW;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static io.prestosql.plugin.bigquery.BigQueryErrorCode.BIGQUERY_FAILED_TO_EXECUTE_QUERY;
+import static io.prestosql.spi.StandardErrorCode.NOT_SUPPORTED;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.IntStream.range;
@@ -124,8 +128,19 @@ private List<BigQuerySplit> createEmptyProjection(TableId tableId, int actualPar
                 numberOfRows = result.iterateAll().iterator().next().get(0).getLongValue();
             }
             else {
-                // no filters, so we can take the value from the table info
-                numberOfRows = bigQueryClient.getTable(tableId).getNumRows().longValue();
+                // no filters, so we can take the value from the table info when the object is TABLE
+                TableInfo tableInfo = bigQueryClient.getTable(tableId);
+                if (tableInfo.getDefinition().getType() == TABLE) {
+                    numberOfRows = tableInfo.getNumRows().longValue();
+                }
+                else if (tableInfo.getDefinition().getType() == VIEW) {
+                    String sql = bigQueryClient.selectSql(tableId, "COUNT(*)");
+                    TableResult result = bigQueryClient.query(sql);
+                    numberOfRows = result.iterateAll().iterator().next().get(0).getLongValue();
+                }
+                else {
+                    throw new PrestoException(NOT_SUPPORTED, "Unsupported table type: " + tableInfo.getDefinition().getType());
+                }
             }
 
             long rowsPerSplit = numberOfRows / actualParallelism;
diff --git a/presto-bigquery/src/test/java/io/prestosql/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java b/presto-bigquery/src/test/java/io/prestosql/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java
@@ -105,6 +105,33 @@ public void testPredicatePushdownPrunnedColumns()
                 "VALUES (1)");
     }
 
+    @Test(description = "regression test for https://github.com/prestosql/presto/issues/5635")
+    public void testCountAggregationView()
+    {
+        BigQuery client = createBigQueryClient();
+
+        String tableName = "test.count_aggregation_table";
+        String viewName = "test.count_aggregation_view";
+
+        executeBigQuerySql(client, "DROP TABLE IF EXISTS " + tableName);
+        executeBigQuerySql(client, "DROP VIEW IF EXISTS " + viewName);
+        executeBigQuerySql(client, "CREATE TABLE " + tableName + " (a INT64, b INT64, c INT64)");
+        executeBigQuerySql(client, "INSERT INTO " + tableName + " VALUES (1, 2, 3), (4, 5, 6)");
+        executeBigQuerySql(client, "CREATE VIEW " + viewName + " AS SELECT * FROM " + tableName);
+
+        assertQuery(
+                "SELECT count(*) FROM " + viewName,
+                "VALUES (2)");
+
+        assertQuery(
+                "SELECT count(*) FROM " + viewName + " WHERE a = 1",
+                "VALUES (1)");
+
+        assertQuery(
+                "SELECT count(a) FROM " + viewName + " WHERE b = 2",
+                "VALUES (1)");
+    }
+
     private static void executeBigQuerySql(BigQuery bigquery, String query)
     {
         QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder(query)
