Implement dereference pushdown for MongoDB connector

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/projection/ApplyProjectionUtil.java

@@ -0,0 +1,172 @@
+/*
+ * 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.base.projection;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
+import io.trino.spi.expression.Call;
+import io.trino.spi.expression.ConnectorExpression;
+import io.trino.spi.expression.Constant;
+import io.trino.spi.expression.FieldDereference;
+import io.trino.spi.expression.Variable;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.function.Predicate;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+import static java.util.Objects.requireNonNull;
+
+public final class ApplyProjectionUtil
+{
+    private ApplyProjectionUtil() {}
+
+    public static List<ConnectorExpression> extractSupportedProjectedColumns(ConnectorExpression expression)
+    {
+        return extractSupportedProjectedColumns(expression, connectorExpression -> true);
+    }
+
+    public static List<ConnectorExpression> extractSupportedProjectedColumns(ConnectorExpression expression, Predicate<ConnectorExpression> expressionPredicate)
+    {
+        requireNonNull(expression, "expression is null");
+        ImmutableList.Builder<ConnectorExpression> supportedSubExpressions = ImmutableList.builder();
+        fillSupportedProjectedColumns(expression, supportedSubExpressions, expressionPredicate);
+        return supportedSubExpressions.build();
+    }
+
+    private static void fillSupportedProjectedColumns(ConnectorExpression expression, ImmutableList.Builder<ConnectorExpression> supportedSubExpressions, Predicate<ConnectorExpression> expressionPredicate)
+    {
+        if (isPushdownSupported(expression, expressionPredicate)) {
+            supportedSubExpressions.add(expression);
+            return;
+        }
+
+        // If the whole expression is not supported, look for a partially supported projection
+        for (ConnectorExpression child : expression.getChildren()) {
+            fillSupportedProjectedColumns(child, supportedSubExpressions, expressionPredicate);
+        }
+    }
+
+    @VisibleForTesting
+    static boolean isPushdownSupported(ConnectorExpression expression, Predicate<ConnectorExpression> expressionPredicate)
+    {
+        return expressionPredicate.test(expression)
+                && (expression instanceof Variable ||
+                    (expression instanceof FieldDereference fieldDereference
+                            && isPushdownSupported(fieldDereference.getTarget(), expressionPredicate)));
+    }
+
+    public static ProjectedColumnRepresentation createProjectedColumnRepresentation(ConnectorExpression expression)
+    {
+        ImmutableList.Builder<Integer> ordinals = ImmutableList.builder();
+
+        Variable target;
+        while (true) {
+            if (expression instanceof Variable variable) {
+                target = variable;
+                break;
+            }
+            if (expression instanceof FieldDereference dereference) {
+                ordinals.add(dereference.getField());
+                expression = dereference.getTarget();
+            }
+            else {
+                throw new IllegalArgumentException("expression is not a valid dereference chain");
+            }
+        }
+
+        return new ProjectedColumnRepresentation(target, ordinals.build().reverse());
+    }
+
+    /**
+     * Replace all connector expressions with variables as given by {@param expressionToVariableMappings} in a top down manner.
+     * i.e. if the replacement occurs for the parent, the children will not be visited.
+     */
+    public static ConnectorExpression replaceWithNewVariables(ConnectorExpression expression, Map<ConnectorExpression, Variable> expressionToVariableMappings)
+    {
+        if (expressionToVariableMappings.containsKey(expression)) {
+            return expressionToVariableMappings.get(expression);
+        }
+
+        if (expression instanceof Constant || expression instanceof Variable) {
+            return expression;
+        }
+
+        if (expression instanceof FieldDereference fieldDereference) {
+            ConnectorExpression newTarget = replaceWithNewVariables(fieldDereference.getTarget(), expressionToVariableMappings);
+            return new FieldDereference(expression.getType(), newTarget, fieldDereference.getField());
+        }
+
+        if (expression instanceof Call call) {
+            return new Call(
+                    call.getType(),
+                    call.getFunctionName(),
+                    call.getArguments().stream()
+                            .map(argument -> replaceWithNewVariables(argument, expressionToVariableMappings))
+                            .collect(toImmutableList()));
+        }
+
+        // We cannot skip processing for unsupported expression shapes. This may lead to variables being left in ProjectionApplicationResult
+        // which are no longer bound.
+        throw new UnsupportedOperationException("Unsupported expression: " + expression);
+    }
+
+    public static class ProjectedColumnRepresentation
+    {
+        private final Variable variable;
+        private final List<Integer> dereferenceIndices;
+
+        public ProjectedColumnRepresentation(Variable variable, List<Integer> dereferenceIndices)
+        {
+            this.variable = requireNonNull(variable, "variable is null");
+            this.dereferenceIndices = ImmutableList.copyOf(requireNonNull(dereferenceIndices, "dereferenceIndices is null"));
+        }
+
+        public Variable getVariable()
+        {
+            return variable;
+        }
+
+        public List<Integer> getDereferenceIndices()
+        {
+            return dereferenceIndices;
+        }
+
+        public boolean isVariable()
+        {
+            return dereferenceIndices.isEmpty();
+        }
+
+        @Override
+        public boolean equals(Object obj)
+        {
+            if (this == obj) {
+                return true;
+            }
+            if ((obj == null) || (getClass() != obj.getClass())) {
+                return false;
+            }
+            ProjectedColumnRepresentation that = (ProjectedColumnRepresentation) obj;
+            return Objects.equals(variable, that.variable) &&
+                    Objects.equals(dereferenceIndices, that.dereferenceIndices);
+        }
+
+        @Override
+        public int hashCode()
+        {
+            return Objects.hash(variable, dereferenceIndices);
+        }
+    }
+}

lib/trino-plugin-toolkit/src/test/java/io/trino/plugin/base/projection/TestApplyProjectionUtil.java

@@ -0,0 +1,126 @@
+/*
+ * 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.base.projection;
+
+import com.google.common.collect.ImmutableList;
+import io.trino.spi.expression.ConnectorExpression;
+import io.trino.spi.expression.Constant;
+import io.trino.spi.expression.FieldDereference;
+import io.trino.spi.expression.Variable;
+import io.trino.spi.type.RowType;
+import org.testng.annotations.Test;
+
+import static io.trino.plugin.base.projection.ApplyProjectionUtil.extractSupportedProjectedColumns;
+import static io.trino.plugin.base.projection.ApplyProjectionUtil.isPushdownSupported;
+import static io.trino.spi.type.IntegerType.INTEGER;
+import static io.trino.spi.type.RowType.field;
+import static io.trino.spi.type.RowType.rowType;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+public class TestApplyProjectionUtil
+{
+    private static final ConnectorExpression ROW_OF_ROW_VARIABLE = new Variable("a", rowType(field("b", rowType(field("c", INTEGER)))));
+    private static final ConnectorExpression LEAF_DOTTED_ROW_OF_ROW_VARIABLE = new Variable("a", rowType(field("b", rowType(field("c.x", INTEGER)))));
+    private static final ConnectorExpression MID_DOTTED_ROW_OF_ROW_VARIABLE = new Variable("a", rowType(field("b.x", rowType(field("c", INTEGER)))));
+
+    private static final ConnectorExpression ONE_LEVEL_DEREFERENCE = new FieldDereference(
+            rowType(field("c", INTEGER)),
+            ROW_OF_ROW_VARIABLE,
+            0);
+
+    private static final ConnectorExpression TWO_LEVEL_DEREFERENCE = new FieldDereference(
+            INTEGER,
+            ONE_LEVEL_DEREFERENCE,
+            0);
+
+    private static final ConnectorExpression LEAF_DOTTED_ONE_LEVEL_DEREFERENCE = new FieldDereference(
+            rowType(field("c.x", INTEGER)),
+            LEAF_DOTTED_ROW_OF_ROW_VARIABLE,
+            0);
+
+    private static final ConnectorExpression LEAF_DOTTED_TWO_LEVEL_DEREFERENCE = new FieldDereference(
+            INTEGER,
+            LEAF_DOTTED_ONE_LEVEL_DEREFERENCE,
+            0);
+
+    private static final ConnectorExpression MID_DOTTED_ONE_LEVEL_DEREFERENCE = new FieldDereference(
+            rowType(field("c.x", INTEGER)),
+            MID_DOTTED_ROW_OF_ROW_VARIABLE,
+            0);
+
+    private static final ConnectorExpression MID_DOTTED_TWO_LEVEL_DEREFERENCE = new FieldDereference(
+            INTEGER,
+            MID_DOTTED_ONE_LEVEL_DEREFERENCE,
+            0);
+
+    private static final ConnectorExpression INT_VARIABLE = new Variable("a", INTEGER);
+    private static final ConnectorExpression CONSTANT = new Constant(5, INTEGER);
+
+    @Test
+    public void testIsProjectionSupported()
+    {
+        assertTrue(isPushdownSupported(ONE_LEVEL_DEREFERENCE, connectorExpression -> true));
+        assertTrue(isPushdownSupported(TWO_LEVEL_DEREFERENCE, connectorExpression -> true));
+        assertTrue(isPushdownSupported(INT_VARIABLE, connectorExpression -> true));
+        assertFalse(isPushdownSupported(CONSTANT, connectorExpression -> true));
+
+        assertFalse(isPushdownSupported(ONE_LEVEL_DEREFERENCE, connectorExpression -> false));
+        assertFalse(isPushdownSupported(TWO_LEVEL_DEREFERENCE, connectorExpression -> false));
+        assertFalse(isPushdownSupported(INT_VARIABLE, connectorExpression -> false));
+        assertFalse(isPushdownSupported(CONSTANT, connectorExpression -> false));
+
+        assertTrue(isPushdownSupported(LEAF_DOTTED_ONE_LEVEL_DEREFERENCE, this::isSupportedForPushDown));
+        assertFalse(isPushdownSupported(LEAF_DOTTED_TWO_LEVEL_DEREFERENCE, this::isSupportedForPushDown));
+        assertFalse(isPushdownSupported(MID_DOTTED_ONE_LEVEL_DEREFERENCE, this::isSupportedForPushDown));
+        assertFalse(isPushdownSupported(MID_DOTTED_TWO_LEVEL_DEREFERENCE, this::isSupportedForPushDown));
+    }
+
+    @Test
+    public void testExtractSupportedProjectionColumns()
+    {
+        assertEquals(extractSupportedProjectedColumns(ONE_LEVEL_DEREFERENCE), ImmutableList.of(ONE_LEVEL_DEREFERENCE));
+        assertEquals(extractSupportedProjectedColumns(TWO_LEVEL_DEREFERENCE), ImmutableList.of(TWO_LEVEL_DEREFERENCE));
+        assertEquals(extractSupportedProjectedColumns(INT_VARIABLE), ImmutableList.of(INT_VARIABLE));
+        assertEquals(extractSupportedProjectedColumns(CONSTANT), ImmutableList.of());
+
+        assertEquals(extractSupportedProjectedColumns(ONE_LEVEL_DEREFERENCE, connectorExpression -> false), ImmutableList.of());
+        assertEquals(extractSupportedProjectedColumns(TWO_LEVEL_DEREFERENCE, connectorExpression -> false), ImmutableList.of());
+        assertEquals(extractSupportedProjectedColumns(INT_VARIABLE, connectorExpression -> false), ImmutableList.of());
+        assertEquals(extractSupportedProjectedColumns(CONSTANT, connectorExpression -> false), ImmutableList.of());
+
+        // Partial supported projection
+        assertEquals(extractSupportedProjectedColumns(LEAF_DOTTED_ONE_LEVEL_DEREFERENCE, this::isSupportedForPushDown), ImmutableList.of(LEAF_DOTTED_ONE_LEVEL_DEREFERENCE));
+        assertEquals(extractSupportedProjectedColumns(LEAF_DOTTED_TWO_LEVEL_DEREFERENCE, this::isSupportedForPushDown), ImmutableList.of(LEAF_DOTTED_ONE_LEVEL_DEREFERENCE));
+        assertEquals(extractSupportedProjectedColumns(MID_DOTTED_ONE_LEVEL_DEREFERENCE, this::isSupportedForPushDown), ImmutableList.of(MID_DOTTED_ROW_OF_ROW_VARIABLE));
+        assertEquals(extractSupportedProjectedColumns(MID_DOTTED_TWO_LEVEL_DEREFERENCE, this::isSupportedForPushDown), ImmutableList.of(MID_DOTTED_ROW_OF_ROW_VARIABLE));
+    }
+
+    /**
+     * This method is used to simulate the behavior when the field passed in the connectorExpression might not supported for pushdown.
+     */
+    private boolean isSupportedForPushDown(ConnectorExpression connectorExpression)
+    {
+        if (connectorExpression instanceof FieldDereference fieldDereference) {
+            RowType rowType = (RowType) fieldDereference.getTarget().getType();
+            RowType.Field field = rowType.getFields().get(fieldDereference.getField());
+            String fieldName = field.getName().get();
+            if (fieldName.contains(".") || fieldName.contains("$")) {
+                return false;
+            }
+        }
+        return true;
+    }
+}

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java

@@ -32,6 +32,7 @@
 import io.trino.filesystem.TrinoFileSystem;
 import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.plugin.base.classloader.ClassLoaderSafeSystemTable;
+import io.trino.plugin.base.projection.ApplyProjectionUtil;
 import io.trino.plugin.deltalake.DeltaLakeAnalyzeProperties.AnalyzeMode;
 import io.trino.plugin.deltalake.expression.ParsingException;
 import io.trino.plugin.deltalake.expression.SparkExpressionParser;
@@ -62,7 +63,6 @@
 import io.trino.plugin.deltalake.transactionlog.writer.TransactionConflictException;
 import io.trino.plugin.deltalake.transactionlog.writer.TransactionLogWriter;
 import io.trino.plugin.deltalake.transactionlog.writer.TransactionLogWriterFactory;
-import io.trino.plugin.hive.HiveApplyProjectionUtil;
 import io.trino.plugin.hive.HiveType;
 import io.trino.plugin.hive.SchemaAlreadyExistsException;
 import io.trino.plugin.hive.TableAlreadyExistsException;
@@ -168,6 +168,9 @@
 import static com.google.common.collect.Sets.difference;
 import static com.google.common.primitives.Ints.max;
 import static io.trino.filesystem.Locations.appendPath;
+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.deltalake.DataFileInfo.DataFileType.DATA;
 import static io.trino.plugin.deltalake.DeltaLakeAnalyzeProperties.AnalyzeMode.FULL_REFRESH;
 import static io.trino.plugin.deltalake.DeltaLakeAnalyzeProperties.AnalyzeMode.INCREMENTAL;
@@ -231,9 +234,6 @@
 import static io.trino.plugin.deltalake.transactionlog.MetadataEntry.configurationForNewTable;
 import static io.trino.plugin.deltalake.transactionlog.TransactionLogParser.getMandatoryCurrentVersion;
 import static io.trino.plugin.deltalake.transactionlog.TransactionLogUtil.getTransactionLogDir;
-import static io.trino.plugin.hive.HiveApplyProjectionUtil.ProjectedColumnRepresentation;
-import static io.trino.plugin.hive.HiveApplyProjectionUtil.extractSupportedProjectedColumns;
-import static io.trino.plugin.hive.HiveApplyProjectionUtil.replaceWithNewVariables;
 import static io.trino.plugin.hive.HiveMetadata.PRESTO_QUERY_ID_NAME;
 import static io.trino.plugin.hive.TableType.EXTERNAL_TABLE;
 import static io.trino.plugin.hive.TableType.MANAGED_TABLE;
@@ -2372,7 +2372,7 @@ public Optional<ProjectionApplicationResult<ConnectorTableHandle>> applyProjecti
                 .collect(toImmutableSet());
 
         Map<ConnectorExpression, ProjectedColumnRepresentation> columnProjections = projectedExpressions.stream()
-                .collect(toImmutableMap(Function.identity(), HiveApplyProjectionUtil::createProjectedColumnRepresentation));
+                .collect(toImmutableMap(Function.identity(), ApplyProjectionUtil::createProjectedColumnRepresentation));
 
         // all references are simple variables
         if (!isProjectionPushdownEnabled(session)