Skip to content

Commit 6a50b47

Browse files
committed
Fix projection pushdown in BigQuery connector
1 parent 8c1afc0 commit 6a50b47

File tree

3 files changed

+33
-1
lines changed

3 files changed

+33
-1
lines changed

presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryMetadata.java

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -267,7 +267,10 @@ public Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProjecti
267267
session, handle, projections, assignments);
268268
BigQueryTableHandle bigQueryTableHandle = (BigQueryTableHandle) handle;
269269

270-
if (bigQueryTableHandle.getProjectedColumns().isPresent()) {
270+
List<ColumnHandle> newColumns = assignments.values().stream()
271+
.collect(toImmutableList());
272+
273+
if (bigQueryTableHandle.getProjectedColumns().isPresent() && containSameElements(newColumns, bigQueryTableHandle.getProjectedColumns().get())) {
271274
return Optional.empty();
272275
}
273276

@@ -303,4 +306,9 @@ public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(
303306

304307
return Optional.of(new ConstraintApplicationResult<>(updatedHandle, constraint.getSummary()));
305308
}
309+
310+
private static boolean containSameElements(Iterable<? extends ColumnHandle> first, Iterable<? extends ColumnHandle> second)
311+
{
312+
return ImmutableSet.copyOf(first).equals(ImmutableSet.copyOf(second));
313+
}
306314
}

presto-bigquery/src/main/java/io/prestosql/plugin/bigquery/BigQueryPageSourceProvider.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727

2828
import java.util.List;
2929

30+
import static com.google.common.base.Preconditions.checkArgument;
3031
import static com.google.common.collect.ImmutableList.toImmutableList;
3132
import static java.util.Objects.requireNonNull;
3233

@@ -56,6 +57,11 @@ public ConnectorPageSource createPageSource(
5657
{
5758
log.debug("createPageSource(transaction=%s, session=%s, split=%s, table=%s, columns=%s)", transaction, session, split, table, columns);
5859
BigQuerySplit bigQuerySplit = (BigQuerySplit) split;
60+
61+
// We expect columns list requested here to match list passed to ConnectorMetadata.applyProjection.
62+
checkArgument(bigQuerySplit.getColumns().isEmpty() || bigQuerySplit.getColumns().equals(columns),
63+
"Requested columns %s do not match list in split %s", columns, bigQuerySplit.getColumns());
64+
5965
if (bigQuerySplit.representsEmptyProjection()) {
6066
return new BigQueryEmptyProjectionPageSource(bigQuerySplit.getEmptyRowsToGenerate());
6167
}

presto-bigquery/src/test/java/io/prestosql/plugin/bigquery/TestBigQueryIntegrationSmokeTest.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,24 @@ public void testSelectFromYearlyPartitionedTable()
8787
assertEquals((long) actualValues.getOnlyValue(), 1L);
8888
}
8989

90+
@Test(description = "regression test for https://github.com/prestosql/presto/issues/5618")
91+
public void testPredicatePushdownPrunnedColumns()
92+
{
93+
BigQuery client = createBigQueryClient();
94+
95+
String tableName = "test.predicate_pushdown_prunned_columns";
96+
97+
executeBigQuerySql(client, "DROP TABLE IF EXISTS " + tableName);
98+
executeBigQuerySql(client, "CREATE TABLE " + tableName + " (a INT64, b INT64, c INT64)");
99+
executeBigQuerySql(client, "INSERT INTO " + tableName + " VALUES (1,2,3)");
100+
101+
assertQuery(
102+
"SELECT 1 FROM " + tableName + " WHERE " +
103+
" ((NULL IS NULL) OR a = 100) AND " +
104+
" b = 2",
105+
"VALUES (1)");
106+
}
107+
90108
private static void executeBigQuerySql(BigQuery bigquery, String query)
91109
{
92110
QueryJobConfiguration queryConfig = QueryJobConfiguration.newBuilder(query)

0 commit comments

Comments
 (0)