Add projection push down for STRUCT field in big query connector

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryArrowToPageConverter.java

@@ -103,16 +103,32 @@ public void convert(PageBuilder pageBuilder, ArrowRecordBatch batch)
 
  for (int column = 0; column < columns.size(); column++) {
  BigQueryColumnHandle columnHandle = columns.get(column);
+ FieldVector fieldVector = getFieldVector(root, columnHandle);
  convertType(pageBuilder.getBlockBuilder(column),
  columnHandle.trinoType(),
- root.getVector(toBigQueryColumnName(columnHandle.name())),
+ fieldVector,
  0,
- root.getVector(toBigQueryColumnName(columnHandle.name())).getValueCount());
+ fieldVector.getValueCount());
  }
 
  root.clear();
  }
 
+ private static FieldVector getFieldVector(VectorSchemaRoot root, BigQueryColumnHandle columnHandle)
+ {
+ FieldVector fieldVector = root.getVector(toBigQueryColumnName(columnHandle.name()));
+
+ for (String dereferenceName : columnHandle.dereferenceNames()) {
+ for (FieldVector child : fieldVector.getChildrenFromFields()) {
+ if (child.getField().getName().equals(dereferenceName)) {
+ fieldVector = child;
+ break;
+ }
+ }
+ }
+ return fieldVector;
+ }
+
  private void convertType(BlockBuilder output, Type type, FieldVector vector, int offset, int length)
  {
  Class<?> javaType = type.getJavaType();

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryClient.java

@@ -33,6 +33,7 @@
 import com.google.cloud.bigquery.TableInfo;
 import com.google.cloud.bigquery.TableResult;
 import com.google.cloud.http.BaseHttpServiceException;
+import com.google.common.base.Joiner;
 import com.google.common.cache.Cache;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
@@ -468,8 +469,17 @@ public TableId getDestinationTable(String sql)
 
  public static String selectSql(TableId table, List<BigQueryColumnHandle> requiredColumns, Optional<String> filter)
  {
- String columns = requiredColumns.stream().map(column -> format("`%s`", column.name())).collect(joining(","));
- return selectSql(table, columns, filter);
+ return selectSql(table,
+ requiredColumns.stream()
+ .map(column -> Joiner.on('.')
+ .join(ImmutableList.<String>builder()
+ .add(format("`%s`", column.name()))
+ .addAll(column.dereferenceNames().stream()
+ .map(dereferenceName -> format("`%s`", dereferenceName))
+ .collect(toImmutableList()))
+ .build()))
+ .collect(joining(",")),
+ filter);
  }
 
  public static String selectSql(TableId table, String formattedColumns, Optional<String> filter)

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryColumnHandle.java

@@ -16,6 +16,7 @@
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.google.cloud.bigquery.Field;
 import com.google.cloud.bigquery.StandardSQLTypeName;
+import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ColumnMetadata;
@@ -30,6 +31,7 @@
 
 public record BigQueryColumnHandle(
  String name,
+ List<String> dereferenceNames,
  Type trinoType,
  StandardSQLTypeName bigqueryType,
  boolean isPushdownSupported,
@@ -44,6 +46,7 @@ public record BigQueryColumnHandle(
  public BigQueryColumnHandle
  {
  requireNonNull(name, "name is null");
+ dereferenceNames = ImmutableList.copyOf(requireNonNull(dereferenceNames, "dereferenceNames is null"));
  requireNonNull(trinoType, "trinoType is null");
  requireNonNull(bigqueryType, "bigqueryType is null");
  requireNonNull(mode, "mode is null");
@@ -62,6 +65,16 @@ public ColumnMetadata getColumnMetadata()
  .build();
  }
 
+ @JsonIgnore
+ public String getQualifiedName()
+ {
+ return Joiner.on('.')
+ .join(ImmutableList.<String>builder()
+ .add(name)
+ .addAll(dereferenceNames)
+ .build());
+ }
+
  @JsonIgnore
  public long getRetainedSizeInBytes()
  {

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryConfig.java

@@ -63,6 +63,7 @@ public class BigQueryConfig
  private String queryLabelName;
  private String queryLabelFormat;
  private boolean proxyEnabled;
+ private boolean projectionPushDownEnabled = true;
  private int metadataParallelism = 2;
 
  public Optional<String> getProjectId()
@@ -342,6 +343,19 @@ public BigQueryConfig setProxyEnabled(boolean proxyEnabled)
  return this;
  }
 
+ public boolean isProjectionPushdownEnabled()
+ {
+ return projectionPushDownEnabled;
+ }
+
+ @Config("bigquery.projection-pushdown-enabled")
+ @ConfigDescription("Dereference push down for ROW type")
+ public BigQueryConfig setProjectionPushdownEnabled(boolean projectionPushDownEnabled)
+ {
+ this.projectionPushDownEnabled = projectionPushDownEnabled;
+ return this;
+ }
+
  @Min(1)
  @Max(32)
  public int getMetadataParallelism()

plugin/trino-bigquery/src/main/java/io/trino/plugin/bigquery/BigQueryMetadata.java

@@ -32,15 +32,18 @@
 import com.google.cloud.bigquery.storage.v1.JsonStreamWriter;
 import com.google.cloud.bigquery.storage.v1.TableName;
 import com.google.cloud.bigquery.storage.v1.WriteStream;
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Functions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Ordering;
 import com.google.common.io.Closer;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import io.airlift.log.Logger;
 import io.airlift.slice.Slice;
+import io.trino.plugin.base.projection.ApplyProjectionUtil;
 import io.trino.plugin.bigquery.BigQueryClient.RemoteDatabaseObject;
 import io.trino.plugin.bigquery.BigQueryTableHandle.BigQueryPartitionType;
 import io.trino.plugin.bigquery.ptf.Query.QueryHandle;
@@ -78,18 +81,22 @@
 import io.trino.spi.connector.TableFunctionApplicationResult;
 import io.trino.spi.connector.TableNotFoundException;
 import io.trino.spi.expression.ConnectorExpression;
+import io.trino.spi.expression.Variable;
 import io.trino.spi.function.table.ConnectorTableFunctionHandle;
 import io.trino.spi.predicate.Domain;
 import io.trino.spi.predicate.TupleDomain;
 import io.trino.spi.security.TrinoPrincipal;
 import io.trino.spi.statistics.ComputedStatistics;
 import io.trino.spi.type.BigintType;
+import io.trino.spi.type.RowType;
 import io.trino.spi.type.Type;
 import io.trino.spi.type.VarcharType;
 import org.json.JSONArray;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
@@ -108,16 +115,22 @@
 import static com.google.cloud.bigquery.storage.v1.WriteStream.Type.COMMITTED;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkState;
+import static com.google.common.base.Verify.verify;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static com.google.common.util.concurrent.Futures.allAsList;
 import static io.trino.plugin.base.TemporaryTables.generateTemporaryTableName;
+import static io.trino.plugin.base.projection.ApplyProjectionUtil.ProjectedColumnRepresentation;
+import static io.trino.plugin.base.projection.ApplyProjectionUtil.extractSupportedProjectedColumns;
+import static io.trino.plugin.base.projection.ApplyProjectionUtil.replaceWithNewVariables;
 import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_BAD_WRITE;
 import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_FAILED_TO_EXECUTE_QUERY;
 import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_LISTING_TABLE_ERROR;
 import static io.trino.plugin.bigquery.BigQueryErrorCode.BIGQUERY_UNSUPPORTED_OPERATION;
 import static io.trino.plugin.bigquery.BigQueryPseudoColumn.PARTITION_DATE;
 import static io.trino.plugin.bigquery.BigQueryPseudoColumn.PARTITION_TIME;
+import static io.trino.plugin.bigquery.BigQuerySessionProperties.isProjectionPushdownEnabled;
 import static io.trino.plugin.bigquery.BigQueryTableHandle.BigQueryPartitionType.INGESTION;
 import static io.trino.plugin.bigquery.BigQueryTableHandle.getPartitionType;
 import static io.trino.plugin.bigquery.BigQueryUtil.isWildcardTable;
@@ -138,6 +151,8 @@ public class BigQueryMetadata
 {
  private static final Logger log = Logger.get(BigQueryMetadata.class);
  private static final Type TRINO_PAGE_SINK_ID_COLUMN_TYPE = BigintType.BIGINT;
+ private static final Ordering<BigQueryColumnHandle> COLUMN_HANDLE_ORDERING = Ordering
+ .from(Comparator.comparingInt(columnHandle -> columnHandle.dereferenceNames().size()));
 
  static final int DEFAULT_NUMERIC_TYPE_PRECISION = 38;
  static final int DEFAULT_NUMERIC_TYPE_SCALE = 9;
@@ -771,7 +786,7 @@ public Optional<ConnectorOutputMetadata> finishInsert(
  @Override
  public ColumnHandle getMergeRowIdColumnHandle(ConnectorSession session, ConnectorTableHandle tableHandle)
  {
- return new BigQueryColumnHandle("$merge_row_id", BIGINT, INT64, true, Field.Mode.REQUIRED, ImmutableList.of(), null, true);
+ return new BigQueryColumnHandle("$merge_row_id", ImmutableList.of(), BIGINT, INT64, true, Field.Mode.REQUIRED, ImmutableList.of(), null, true);
  }
 
  @Override
@@ -882,24 +897,150 @@ public Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProjecti
  log.debug("applyProjection(session=%s, handle=%s, projections=%s, assignments=%s)",
  session, handle, projections, assignments);
  BigQueryTableHandle bigQueryTableHandle = (BigQueryTableHandle) handle;
+ if (!isProjectionPushdownEnabled(session)) {
+ List<ColumnHandle> newColumns = ImmutableList.copyOf(assignments.values());
+ if (bigQueryTableHandle.projectedColumns().isPresent() && containSameElements(newColumns, bigQueryTableHandle.projectedColumns().get())) {
+ return Optional.empty();
+ }
 
- List<ColumnHandle> newColumns = ImmutableList.copyOf(assignments.values());
+ ImmutableList.Builder<BigQueryColumnHandle> projectedColumns = ImmutableList.builder();
+ ImmutableList.Builder<Assignment> assignmentList = ImmutableList.builder();
+ assignments.forEach((name, column) -> {
+ BigQueryColumnHandle columnHandle = (BigQueryColumnHandle) column;
+ projectedColumns.add(columnHandle);
+ assignmentList.add(new Assignment(name, column, columnHandle.trinoType()));
+ });
 
- if (bigQueryTableHandle.projectedColumns().isPresent() && containSameElements(newColumns, bigQueryTableHandle.projectedColumns().get())) {
- return Optional.empty();
+ bigQueryTableHandle = bigQueryTableHandle.withProjectedColumns(projectedColumns.build());
+
+ return Optional.of(new ProjectionApplicationResult<>(bigQueryTableHandle, projections, assignmentList.build(), false));
  }
 
- ImmutableList.Builder<BigQueryColumnHandle> projectedColumns = ImmutableList.builder();
- ImmutableList.Builder<Assignment> assignmentList = ImmutableList.builder();
- assignments.forEach((name, column) -> {
- BigQueryColumnHandle columnHandle = (BigQueryColumnHandle) column;
- projectedColumns.add(columnHandle);
- assignmentList.add(new Assignment(name, column, columnHandle.trinoType()));
- });
+ // Create projected column representations for supported sub expressions. Simple column references and chain of
+ // dereferences on a variable are supported right now.
+ Set<ConnectorExpression> projectedExpressions = projections.stream()
+ .flatMap(expression -> extractSupportedProjectedColumns(expression).stream())
+ .collect(toImmutableSet());
+
+ Map<ConnectorExpression, ProjectedColumnRepresentation> columnProjections = projectedExpressions.stream()
+ .collect(toImmutableMap(identity(), ApplyProjectionUtil::createProjectedColumnRepresentation));
+
+ // all references are simple variables
+ if (columnProjections.values().stream().allMatch(ProjectedColumnRepresentation::isVariable)) {
+ Set<BigQueryColumnHandle> projectedColumns = ImmutableSet.copyOf(projectParentColumns(assignments.values().stream()
+ .map(BigQueryColumnHandle.class::cast)
+ .collect(toImmutableList())));
+ if (bigQueryTableHandle.projectedColumns().isPresent() && containSameElements(projectedColumns, bigQueryTableHandle.projectedColumns().get())) {
+ return Optional.empty();
+ }
+ List<Assignment> assignmentsList = assignments.entrySet().stream()
+ .map(assignment -> new Assignment(
+ assignment.getKey(),
+ assignment.getValue(),
+ ((BigQueryColumnHandle) assignment.getValue()).trinoType()))
+ .collect(toImmutableList());
+
+ return Optional.of(new ProjectionApplicationResult<>(
+ bigQueryTableHandle.withProjectedColumns(ImmutableList.copyOf(projectedColumns)),
+ projections,
+ assignmentsList,
+ false));
+ }
+
+ Map<String, Assignment> newAssignments = new HashMap<>();
+ ImmutableMap.Builder<ConnectorExpression, Variable> newVariablesBuilder = ImmutableMap.builder();
+ ImmutableSet.Builder<BigQueryColumnHandle> projectedColumnsBuilder = ImmutableSet.builder();
+
+ for (Map.Entry<ConnectorExpression, ProjectedColumnRepresentation> entry : columnProjections.entrySet()) {
+ ConnectorExpression expression = entry.getKey();
+ ProjectedColumnRepresentation projectedColumn = entry.getValue();
 
- bigQueryTableHandle = bigQueryTableHandle.withProjectedColumns(projectedColumns.build());
+ BigQueryColumnHandle baseColumnHandle = (BigQueryColumnHandle) assignments.get(projectedColumn.getVariable().getName());
+ BigQueryColumnHandle projectedColumnHandle = createProjectedColumnHandle(baseColumnHandle, projectedColumn.getDereferenceIndices(), expression.getType());
+ String projectedColumnName = projectedColumnHandle.getQualifiedName();
+
+ Variable projectedColumnVariable = new Variable(projectedColumnName, expression.getType());
+ Assignment newAssignment = new Assignment(projectedColumnName, projectedColumnHandle, expression.getType());
+ newAssignments.putIfAbsent(projectedColumnName, newAssignment);
+
+ newVariablesBuilder.put(expression, projectedColumnVariable);
+ projectedColumnsBuilder.add(projectedColumnHandle);
+ }
+
+ // Modify projections to refer to new variables
+ Map<ConnectorExpression, Variable> newVariables = newVariablesBuilder.buildOrThrow();
+ List<ConnectorExpression> newProjections = projections.stream()
+ .map(expression -> replaceWithNewVariables(expression, newVariables))
+ .collect(toImmutableList());
+
+ List<Assignment> outputAssignments = newAssignments.values().stream().collect(toImmutableList());
+ return Optional.of(new ProjectionApplicationResult<>(
+ bigQueryTableHandle.withProjectedColumns(projectParentColumns(ImmutableList.copyOf(projectedColumnsBuilder.build()))),
+ newProjections,
+ outputAssignments,
+ false));
+ }
 
- return Optional.of(new ProjectionApplicationResult<>(bigQueryTableHandle, projections, assignmentList.build(), false));
+ /**
+ * Creates a set of parent columns for the input projected columns. For example,
+ * if input {@param columns} include columns "a.b" and "a.b.c", then they will be projected from a single column "a.b".
+ */
+ @VisibleForTesting
+ static List<BigQueryColumnHandle> projectParentColumns(List<BigQueryColumnHandle> columnHandles)
+ {
+ List<BigQueryColumnHandle> sortedColumnHandles = COLUMN_HANDLE_ORDERING.sortedCopy(columnHandles);
+ List<BigQueryColumnHandle> parentColumns = new ArrayList<>();
+ for (BigQueryColumnHandle column : sortedColumnHandles) {
+ if (!parentColumnExists(parentColumns, column)) {
+ parentColumns.add(column);
+ }
+ }
+ return parentColumns;
+ }
+
+ private static boolean parentColumnExists(List<BigQueryColumnHandle> existingColumns, BigQueryColumnHandle column)
+ {
+ for (BigQueryColumnHandle existingColumn : existingColumns) {
+ List<String> existingColumnDereferenceNames = existingColumn.dereferenceNames();
+ verify(
+ column.dereferenceNames().size() >= existingColumnDereferenceNames.size(),
+ "Selected column's dereference size must be greater than or equal to the existing column's dereference size");
+ if (existingColumn.name().equals(column.name())
+ && column.dereferenceNames().subList(0, existingColumnDereferenceNames.size()).equals(existingColumnDereferenceNames)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ private BigQueryColumnHandle createProjectedColumnHandle(BigQueryColumnHandle baseColumn, List<Integer> indices, Type projectedColumnType)
+ {
+ if (indices.isEmpty()) {
+ return baseColumn;
+ }
+
+ ImmutableList.Builder<String> dereferenceNamesBuilder = ImmutableList.builder();
+ dereferenceNamesBuilder.addAll(baseColumn.dereferenceNames());
+
+ Type type = baseColumn.trinoType();
+ for (int index : indices) {
+ checkArgument(type instanceof RowType, "type should be Row type");
+ RowType rowType = (RowType) type;
+ RowType.Field field = rowType.getFields().get(index);
+ dereferenceNamesBuilder.add(field.getName()
+ .orElseThrow(() -> new TrinoException(NOT_SUPPORTED, "ROW type does not have field names declared: " + rowType)));
+ type = field.getType();
+ }
+ return new BigQueryColumnHandle(
+ baseColumn.name(),
+ dereferenceNamesBuilder.build(),
+ projectedColumnType,
+ typeManager.toStandardSqlTypeName(projectedColumnType),
+ baseColumn.isPushdownSupported(),
+ baseColumn.mode(),
+ baseColumn.subColumns(),
+ baseColumn.description(),
+ baseColumn.hidden());
  }
 
  @Override