Fix projection pushdown in BigQuery connector

trinodb · Oct 26, 2020 · 6a50b47 · 6a50b47
1 parent 8c1afc0
commit 6a50b47
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 1 deletion.
diff --git a/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryMetadata.java b/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryMetadata.java
@@ -267,7 +267,10 @@ public Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProjecti
  session, handle, projections, assignments);
  BigQueryTableHandle bigQueryTableHandle = (BigQueryTableHandle) handle;
 
- if (bigQueryTableHandle.getProjectedColumns().isPresent()) {
+ List<ColumnHandle> newColumns = assignments.values().stream()
+ .collect(toImmutableList());
+
+ if (bigQueryTableHandle.getProjectedColumns().isPresent() && containSameElements(newColumns, bigQueryTableHandle.getProjectedColumns().get())) {
  return Optional.empty();
  }
 
@@ -303,4 +306,9 @@ public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(
 
  return Optional.of(new ConstraintApplicationResult<>(updatedHandle, constraint.getSummary()));
  }
+
+ private static boolean containSameElements(Iterable<? extends ColumnHandle> first, Iterable<? extends ColumnHandle> second)
+ {
+ return ImmutableSet.copyOf(first).equals(ImmutableSet.copyOf(second));
+ }
 }
diff --git a/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryPageSourceProvider.java b/presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryPageSourceProvider.java
@@ -27,6 +27,7 @@
 
 import java.util.List;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.util.Objects.requireNonNull;
 
@@ -56,6 +57,11 @@ public ConnectorPageSource createPageSource(
  {
  log.debug("createPageSource(transaction=%s, session=%s, split=%s, table=%s, columns=%s)", transaction, session, split, table, columns);
  BigQuerySplit bigQuerySplit = (BigQuerySplit) split;
+
+ // We expect columns list requested here to match list passed to ConnectorMetadata.applyProjection.
+ checkArgument(bigQuerySplit.getColumns().isEmpty() || bigQuerySplit.getColumns().equals(columns),
+ "Requested columns %s do not match list in split %s", columns, bigQuerySplit.getColumns());
+
  if (bigQuerySplit.representsEmptyProjection()) {
  return new BigQueryEmptyProjectionPageSource(bigQuerySplit.getEmptyRowsToGenerate());
  }

diff --git a/...bigquery/src/test/java/io/prestosql/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java b/...bigquery/src/test/java/io/prestosql/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java
@@ -87,6 +87,24 @@ public void testSelectFromYearlyPartitionedTable()
  assertEquals((long) actualValues.getOnlyValue(), 1L);
  }
 
+ @Test(description = "regression test for https://github.com/prestosql/presto/issues/5618")
+ public void testPredicatePushdownPrunnedColumns()
+ {
+ BigQuery client = createBigQueryClient();
+
+ String tableName = "test.predicate_pushdown_prunned_columns";
+
+ executeBigQuerySql(client, "DROP TABLE IF EXISTS " + tableName);
+ executeBigQuerySql(client, "CREATE TABLE " + tableName + " (a INT64, b INT64, c INT64)");
+ executeBigQuerySql(client, "INSERT INTO " + tableName + " VALUES (1,2,3)");
+
+ assertQuery(
+ "SELECT 1 FROM " + tableName + " WHERE " +
+ " ((NULL IS NULL) OR a = 100) AND " +
+ " b = 2",
+ "VALUES (1)");
+ }
+
  private static void executeBigQuerySql(BigQuery bigquery, String query)
  {
  QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder(query)