Add projection pushdown support to MongoDB

plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoMetadata.java

@@ -15,6 +15,7 @@
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Streams;
 import com.google.common.io.Closer;
 import com.mongodb.client.MongoCollection;
@@ -23,6 +24,7 @@
 import io.trino.plugin.mongodb.MongoIndex.MongodbIndexKey;
 import io.trino.plugin.mongodb.ptf.Query.QueryFunctionHandle;
 import io.trino.spi.TrinoException;
+import io.trino.spi.connector.Assignment;
 import io.trino.spi.connector.ColumnHandle;
 import io.trino.spi.connector.ColumnMetadata;
 import io.trino.spi.connector.ConnectorInsertTableHandle;
@@ -39,12 +41,14 @@
 import io.trino.spi.connector.LimitApplicationResult;
 import io.trino.spi.connector.LocalProperty;
 import io.trino.spi.connector.NotFoundException;
+import io.trino.spi.connector.ProjectionApplicationResult;
 import io.trino.spi.connector.RetryMode;
 import io.trino.spi.connector.SchemaTableName;
 import io.trino.spi.connector.SchemaTablePrefix;
 import io.trino.spi.connector.SortingProperty;
 import io.trino.spi.connector.TableFunctionApplicationResult;
 import io.trino.spi.connector.TableNotFoundException;
+import io.trino.spi.expression.ConnectorExpression;
 import io.trino.spi.predicate.Domain;
 import io.trino.spi.predicate.TupleDomain;
 import io.trino.spi.ptf.ConnectorTableFunctionHandle;
@@ -552,7 +556,7 @@ public Optional<LimitApplicationResult<ConnectorTableHandle>> applyLimit(Connect
         }
 
         return Optional.of(new LimitApplicationResult<>(
-                new MongoTableHandle(handle.getSchemaTableName(), handle.getRemoteTableName(), handle.getFilter(), handle.getConstraint(), OptionalInt.of(toIntExact(limit))),
+                new MongoTableHandle(handle.getSchemaTableName(), handle.getRemoteTableName(), handle.getFilter(), handle.getConstraint(), handle.getProjectedColumns(), OptionalInt.of(toIntExact(limit))),
                 true,
                 false));
     }
@@ -599,11 +603,41 @@ public Optional<ConstraintApplicationResult<ConnectorTableHandle>> applyFilter(C
                 handle.getRemoteTableName(),
                 handle.getFilter(),
                 newDomain,
+                handle.getProjectedColumns(),
                 handle.getLimit());
 
         return Optional.of(new ConstraintApplicationResult<>(handle, remainingFilter, false));
     }
 
+    @Override
+    public Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProjection(
+            ConnectorSession session,
+            ConnectorTableHandle handle,
+            List<ConnectorExpression> projections,
+            Map<String, ColumnHandle> assignments)
+    {
+        MongoTableHandle mongoTableHandle = (MongoTableHandle) handle;
+
+        Set<ColumnHandle> newColumns = ImmutableSet.copyOf(assignments.values());
+
+        if (newColumns.equals(mongoTableHandle.getProjectedColumns())) {
+            return Optional.empty();
+        }
+
+        // TODO: support dereference pushdown
+        ImmutableSet.Builder<MongoColumnHandle> projectedColumns = ImmutableSet.builder();
+        ImmutableList.Builder<Assignment> assignmentList = ImmutableList.builder();
+        assignments.forEach((name, column) -> {
+            MongoColumnHandle columnHandle = (MongoColumnHandle) column;
+            assignmentList.add(new Assignment(name, column, columnHandle.getType()));
+            projectedColumns.add(columnHandle);
+        });
+
+        mongoTableHandle = mongoTableHandle.withProjectedColumns(projectedColumns.build());
+
+        return Optional.of(new ProjectionApplicationResult<>(mongoTableHandle, projections, assignmentList.build(), false));
+    }
+
     @Override
     public Optional<TableFunctionApplicationResult<ConnectorTableHandle>> applyTableFunction(ConnectorSession session, ConnectorTableFunctionHandle handle)
     {

plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java

@@ -455,8 +455,16 @@ public long deleteDocuments(RemoteTableName remoteTableName, TupleDomain<ColumnH
     public MongoCursor<Document> execute(MongoTableHandle tableHandle, List<MongoColumnHandle> columns)
     {
         Document output = new Document();
-        for (MongoColumnHandle column : columns) {
-            output.append(column.getName(), 1);
+
+        if (!tableHandle.getProjectedColumns().isEmpty()) {
+            for (MongoColumnHandle column : tableHandle.getProjectedColumns()) {
+                output.append(column.getName(), 1);
+            }
+        }
+        else {
+            for (MongoColumnHandle column : columns) {
+                output.append(column.getName(), 1);
+            }
         }
         MongoCollection<Document> collection = getCollection(tableHandle.getRemoteTableName());
         Document filter = buildFilter(tableHandle);

plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoTableHandle.java

@@ -23,6 +23,7 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.OptionalInt;
+import java.util.Set;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static java.util.Objects.requireNonNull;
@@ -33,12 +34,13 @@ public class MongoTableHandle
     private final SchemaTableName schemaTableName;
     private final RemoteTableName remoteTableName;
     private final TupleDomain<ColumnHandle> constraint;
+    private final Set<MongoColumnHandle> projectedColumns;
     private final Optional<String> filter;
     private final OptionalInt limit;
 
     public MongoTableHandle(SchemaTableName schemaTableName, RemoteTableName remoteTableName, Optional<String> filter)
     {
-        this(schemaTableName, remoteTableName, filter, TupleDomain.all(), OptionalInt.empty());
+        this(schemaTableName, remoteTableName, filter, TupleDomain.all(), Set.of(), OptionalInt.empty());
     }
 
     @JsonCreator
@@ -47,12 +49,14 @@ public MongoTableHandle(
             @JsonProperty("remoteTableName") RemoteTableName remoteTableName,
             @JsonProperty("filter") Optional<String> filter,
             @JsonProperty("constraint") TupleDomain<ColumnHandle> constraint,
+            @JsonProperty("projectedColumns") Set<MongoColumnHandle> projectedColumns,
             @JsonProperty("limit") OptionalInt limit)
     {
         this.schemaTableName = requireNonNull(schemaTableName, "schemaTableName is null");
         this.remoteTableName = requireNonNull(remoteTableName, "remoteTableName is null");
         this.filter = requireNonNull(filter, "filter is null");
         this.constraint = requireNonNull(constraint, "constraint is null");
+        this.projectedColumns = requireNonNull(projectedColumns, "projectedColumns is null");
         this.limit = requireNonNull(limit, "limit is null");
     }
 
@@ -86,10 +90,16 @@ public OptionalInt getLimit()
         return limit;
     }
 
+    @JsonProperty
+    public Set<MongoColumnHandle> getProjectedColumns()
+    {
+        return projectedColumns;
+    }
+
     @Override
     public int hashCode()
     {
-        return Objects.hash(schemaTableName, filter, constraint, limit);
+        return Objects.hash(schemaTableName, filter, constraint, projectedColumns, limit);
     }
 
     @Override
@@ -106,6 +116,7 @@ public boolean equals(Object obj)
                 Objects.equals(this.remoteTableName, other.remoteTableName) &&
                 Objects.equals(this.filter, other.filter) &&
                 Objects.equals(this.constraint, other.constraint) &&
+                Objects.equals(this.projectedColumns, other.projectedColumns) &&
                 Objects.equals(this.limit, other.limit);
     }
 
@@ -117,7 +128,19 @@ public String toString()
                 .add("remoteTableName", remoteTableName)
                 .add("filter", filter)
                 .add("limit", limit)
+                .add("projectedColumns", projectedColumns)
                 .add("constraint", constraint)
                 .toString();
     }
+
+    public MongoTableHandle withProjectedColumns(Set<MongoColumnHandle> projectedColumns)
+    {
+        return new MongoTableHandle(
+                schemaTableName,
+                remoteTableName,
+                filter,
+                constraint,
+                projectedColumns,
+                limit);
+    }
 }

plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java

@@ -143,6 +143,12 @@ public void testSortItemsReflectedInExplain()
                 "TopNPartial\\[count = 5, orderBy = \\[nationkey DESC");
     }
 
+    @Test
+    public void testProjectionPushdown()
+    {
+        assertThat(query("SELECT orderdate, clerk FROM orders")).isFullyPushedDown();
+    }
+
     @Override
     protected Optional<DataMappingTestSetup> filterDataMappingSmokeTestData(DataMappingTestSetup dataMappingTestSetup)
     {