diff --git a/plugin/trino-bigquery/pom.xml b/plugin/trino-bigquery/pom.xml index d12e346bcb82..64198c00f39e 100644 --- a/plugin/trino-bigquery/pom.xml +++ b/plugin/trino-bigquery/pom.xml @@ -517,6 +517,7 @@ **/TestBigQueryMetadataCaching.java **/TestBigQueryAvroTypeMapping.java **/TestBigQueryArrowTypeMapping.java + **/TestBigQueryArrowSerialization.java **/TestBigQueryMetadata.java **/TestBigQueryInstanceCleaner.java **/TestBigQueryCaseInsensitiveMapping.java @@ -590,6 +591,7 @@ **/TestBigQueryWithDifferentProjectIdConnectorSmokeTest.java **/TestBigQueryWithProxyTest.java + **/TestBigQueryArrowSerialization.java **/TestBigQueryArrowTypeMapping.java **/TestBigQueryAvroTypeMapping.java **/TestBigQueryMetadataCaching.java diff --git a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java index 6b3b9450afd8..0ff1d75ff173 100644 --- a/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java +++ b/plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java @@ -103,7 +103,6 @@ public BigQueryArrowToPageConverter(BigQueryTypeManager typeManager, BufferAlloc .map(field -> field.createVector(allocator)) .collect(toImmutableList()); root = new VectorSchemaRoot(vectors); - verify(vectors.size() == columns.size(), "Vectors, columns size differ"); loader = new VectorLoader(root, INSTANCE); } diff --git a/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryArrowSerialization.java b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryArrowSerialization.java new file mode 100644 index 000000000000..abdda7e74adb --- /dev/null +++ b/plugin/trino-bigquery/src/test/java/io/trino/plugin/bigquery/TestBigQueryArrowSerialization.java @@ -0,0 +1,67 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.trino.plugin.bigquery; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import io.trino.plugin.bigquery.BigQueryQueryRunner.BigQuerySqlExecutor; +import io.trino.testing.AbstractTestQueryFramework; +import io.trino.testing.QueryRunner; +import org.intellij.lang.annotations.Language; +import org.junit.jupiter.api.Test; + +import static io.trino.testing.TestingNames.randomNameSuffix; + +public class TestBigQueryArrowSerialization + extends AbstractTestQueryFramework +{ + private final BigQuerySqlExecutor bigQuerySqlExecutor; + + public TestBigQueryArrowSerialization() + { + this.bigQuerySqlExecutor = new BigQuerySqlExecutor(); + } + + @Override + protected QueryRunner createQueryRunner() + throws Exception + { + return BigQueryQueryRunner.createQueryRunner( + ImmutableMap.of(), + ImmutableMap.of("bigquery.experimental.arrow-serialization.enabled", "true"), + ImmutableList.of()); + } + + @Test + void testViewProjectionPredicates() + { + // Run query that columns in WHERE clause don't exist in projections. + // Columns in ReadSession.setRowRestriction in ReadSessionCreator must exist in selected fields. + String viewName = "test_projection_predicates_" + randomNameSuffix(); + + onBigQuery("CREATE VIEW test." + viewName + " AS SELECT 1 AS id, 'test' AS data"); + try { + assertQuery("SELECT data FROM test." + viewName + " WHERE id = 1", "VALUES 'test'"); + assertQuery("SELECT id FROM test." + viewName + " WHERE data = 'test'", "VALUES 1"); + } + finally { + onBigQuery("DROP VIEW test." + viewName); + } + } + + private void onBigQuery(@Language("SQL") String sql) + { + bigQuerySqlExecutor.execute(sql); + } +}