Fix reading mixed-type array in Mongodb

chenjian2664 · ebyhr · commit 76808e820a97 · 2025-09-10T16:22:26.000+09:00
diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoPageSource.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoPageSource.java
@@ -66,6 +66,7 @@
 import static io.trino.plugin.mongodb.MongoSession.DATABASE_NAME;
 import static io.trino.plugin.mongodb.MongoSession.ID;
 import static io.trino.plugin.mongodb.ObjectIdType.OBJECT_ID;
+import static io.trino.plugin.mongodb.TypeUtils.getImplicitRowFieldIndex;
 import static io.trino.plugin.mongodb.TypeUtils.isJsonType;
 import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
 import static io.trino.spi.type.BigintType.BIGINT;
@@ -103,18 +104,21 @@ public class MongoPageSource
     private boolean finished;
 
     private final PageBuilder pageBuilder;
+    private final String implicitPrefix;
 
     public MongoPageSource(
             MongoSession mongoSession,
             MongoTableHandle tableHandle,
-            List<MongoColumnHandle> columns)
+            List<MongoColumnHandle> columns,
+            String implicitPrefix)
     {
         this.columns = ImmutableList.copyOf(requireNonNull(columns, "columns is null"));
         this.columnTypes = columns.stream().map(MongoColumnHandle::type).collect(toList());
         this.cursor = mongoSession.execute(tableHandle, columns);
         currentDoc = null;
 
         pageBuilder = new PageBuilder(columnTypes);
+        this.implicitPrefix = requireNonNull(implicitPrefix, "implicitPrefix is null");
     }
 
     @Override
@@ -383,7 +387,7 @@ else if (type instanceof RowType rowType) {
         output.appendNull();
     }
 
-    private static Object getColumnValue(Document document, MongoColumnHandle mongoColumnHandle)
+    private Object getColumnValue(Document document, MongoColumnHandle mongoColumnHandle)
     {
         Object value = document.get(mongoColumnHandle.baseName());
         if (mongoColumnHandle.isBaseColumn()) {
@@ -392,17 +396,27 @@ private static Object getColumnValue(Document document, MongoColumnHandle mongoC
         if (value instanceof DBRef dbRefValue) {
             return getDbRefValue(dbRefValue, mongoColumnHandle);
         }
-        Document documentValue = (Document) value;
+
         for (String dereferenceName : mongoColumnHandle.dereferenceNames()) {
             // When parent field itself is null
-            if (documentValue == null) {
+            if (value == null) {
                 return null;
             }
-            value = documentValue.get(dereferenceName);
-            if (value instanceof Document nestedDocument) {
-                documentValue = nestedDocument;
+            if (value instanceof Document documentValue) {
+                value = documentValue.get(dereferenceName);
+            }
+            else {
+                checkArgument(value instanceof List<?>, "Unsupported dereference of %s in %s", value.getClass(), mongoColumnHandle);
+                List<?> arrayValue = (List<?>) value;
+                int arrayPosition = getImplicitRowFieldIndex(dereferenceName, implicitPrefix) - 1;
+                checkArgument(arrayPosition >= 0, "Invalid array position %s in %s", dereferenceName, mongoColumnHandle);
+                if (arrayPosition >= arrayValue.size()) {
+                    return null;
+                }
+                value = arrayValue.get(arrayPosition);
             }
-            else if (value instanceof DBRef dbRefValue) {
+
+            if (value instanceof DBRef dbRefValue) {
                 // Assuming DBRefField is the leaf field
                 return getDbRefValue(dbRefValue, mongoColumnHandle);
             }
diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoPageSourceProvider.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoPageSourceProvider.java
@@ -37,11 +37,13 @@ public class MongoPageSourceProvider
     private static final int MONGO_DOMAIN_COMPACTION_THRESHOLD = 1000;
 
     private final MongoSession mongoSession;
+    private final String implicitPrefix;
 
     @Inject
-    public MongoPageSourceProvider(MongoSession mongoSession)
+    public MongoPageSourceProvider(MongoSession mongoSession, MongoClientConfig config)
     {
         this.mongoSession = requireNonNull(mongoSession, "mongoSession is null");
+        this.implicitPrefix = config.getImplicitRowFieldPrefix();
     }
 
     @Override
@@ -80,6 +82,6 @@ public ConnectorPageSource createPageSource(
             return new EmptyPageSource();
         }
 
-        return new MongoPageSource(mongoSession, newTableHandle, handles.build());
+        return new MongoPageSource(mongoSession, newTableHandle, handles.build(), implicitPrefix);
     }
 }
diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java
@@ -524,7 +524,7 @@ public MongoCursor<Document> execute(MongoTableHandle tableHandle, List<MongoCol
         Set<MongoColumnHandle> projectedColumns = tableHandle.projectedColumns();
         checkArgument(projectedColumns.isEmpty() || projectedColumns.containsAll(columns), "projectedColumns must be empty or equal to columns");
 
-        Document projection = buildProjection(columns);
+        Document projection = buildProjection(columns, implicitPrefix);
 
         MongoCollection<Document> collection = getCollection(tableHandle.remoteTableName());
         Document filter = buildFilter(tableHandle);
@@ -540,7 +540,7 @@ public MongoCursor<Document> execute(MongoTableHandle tableHandle, List<MongoCol
     }
 
     @VisibleForTesting
-    static Document buildProjection(List<MongoColumnHandle> columns)
+    static Document buildProjection(List<MongoColumnHandle> columns, String implicitPrefix)
     {
         Document output = new Document();
 
@@ -552,7 +552,13 @@ static Document buildProjection(List<MongoColumnHandle> columns)
         // Starting in MongoDB 4.4, it is illegal to project an embedded document with any of the embedded document's fields
         // (https://www.mongodb.com/docs/manual/reference/limits/#mongodb-limit-Projection-Restrictions). So, Project only sufficient columns.
         for (MongoColumnHandle column : projectSufficientColumns(columns)) {
-            output.append(column.getQualifiedName(), 1);
+            if (column.dereferenceNames().stream().anyMatch(columnName -> TypeUtils.isImplicitRowField(columnName, implicitPrefix))) {
+                // Add parent field for implicit column
+                output.append(column.baseName(), 1);
+            }
+            else {
+                output.append(column.getQualifiedName(), 1);
+            }
         }
 
         return output;
diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/TypeUtils.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/TypeUtils.java
@@ -21,6 +21,7 @@
 
 import java.util.Set;
 
+import static com.google.common.base.Verify.verify;
 import static io.trino.spi.type.BigintType.BIGINT;
 import static io.trino.spi.type.BooleanType.BOOLEAN;
 import static io.trino.spi.type.DateType.DATE;
@@ -64,4 +65,25 @@ public static boolean isPushdownSupportedType(Type type)
                 || type instanceof ObjectIdType
                 || PUSHDOWN_SUPPORTED_PRIMITIVE_TYPES.contains(type);
     }
+
+    public static boolean isImplicitRowField(String fieldName, String implicitPrefix)
+    {
+        return getImplicitRowFieldIndex(fieldName, implicitPrefix) != -1;
+    }
+
+    public static int getImplicitRowFieldIndex(String fieldName, String implicitPrefix)
+    {
+        if (fieldName.length() <= implicitPrefix.length() || !fieldName.startsWith(implicitPrefix)) {
+            return -1;
+        }
+
+        try {
+            int fieldIndex = Integer.parseInt(fieldName.substring(implicitPrefix.length()));
+            verify(fieldIndex >= 1, "Field index must be >= 1");
+            return fieldIndex;
+        }
+        catch (NumberFormatException e) {
+            return -1;
+        }
+    }
 }
diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java
@@ -123,6 +123,30 @@ protected TestTable createTableWithDefaultColumns()
         return abort("MongoDB connector does not support column default values");
     }
 
+    @Test
+    void testMongoMixedTypeArrayType()
+    {
+        String schema = getSession().getSchema().orElseThrow();
+        String table = "test_mixed_array_" + randomNameSuffix();
+        MongoDatabase db = client.getDatabase(schema);
+
+        db.createCollection(table);
+        db.getCollection(table)
+                .insertOne(new Document("mixed_array_col", ImmutableList.of(1, "two", 3.0, new Document("nested_arr", ImmutableList.of(4, 5)))));
+
+        assertThat(query("SHOW COLUMNS FROM " + table))
+                .skippingTypesCheck()
+                .matches("VALUES " +
+                         "('mixed_array_col', 'row(_pos1 bigint, _pos2 varchar, _pos3 double, _pos4 row(nested_arr array(bigint)))', '', '')");
+
+        assertThat(query("SELECT mixed_array_col._pos1, mixed_array_col._pos2, mixed_array_col._pos3 FROM " + table))
+                .matches("VALUES (BIGINT '1', VARCHAR 'two', DOUBLE '3.0')");
+        assertThat(query("SELECT mixed_array_col._pos4.nested_arr[1], mixed_array_col._pos4.nested_arr[2] FROM " + table))
+                .matches("VALUES (BIGINT '4', BIGINT '5')");
+
+        assertUpdate("DROP TABLE " + table);
+    }
+
     @Test
     @Override
     public void testColumnName()
diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoSession.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoSession.java
@@ -55,7 +55,7 @@ public void testBuildProjectionWithoutId()
     {
         List<MongoColumnHandle> columns = ImmutableList.of(COL1, COL2);
 
-        Document output = MongoSession.buildProjection(columns);
+        Document output = MongoSession.buildProjection(columns, "_pos");
         Document expected = new Document()
                 .append(COL1.baseName(), 1)
                 .append(COL2.baseName(), 1)
@@ -69,7 +69,7 @@ public void testBuildProjectionWithId()
     {
         List<MongoColumnHandle> columns = ImmutableList.of(COL1, COL2, ID_COL);
 
-        Document output = MongoSession.buildProjection(columns);
+        Document output = MongoSession.buildProjection(columns, "_pos");
         Document expected = new Document()
                 .append(COL1.baseName(), 1)
                 .append(COL2.baseName(), 1)

Original file line number	Diff line number	Diff line change
`@@ -37,11 +37,13 @@ public class MongoPageSourceProvider`
`37`	`37`	`private static final int MONGO_DOMAIN_COMPACTION_THRESHOLD = 1000;`
`38`	`38`
`39`	`39`	`private final MongoSession mongoSession;`
	`40`	`+ private final String implicitPrefix;`
`40`	`41`
`41`	`42`	`@Inject`
`42`		`- public MongoPageSourceProvider(MongoSession mongoSession)`
	`43`	`+ public MongoPageSourceProvider(MongoSession mongoSession, MongoClientConfig config)`
`43`	`44`	`{`
`44`	`45`	`this.mongoSession = requireNonNull(mongoSession, "mongoSession is null");`
	`46`	`+ this.implicitPrefix = config.getImplicitRowFieldPrefix();`
`45`	`47`	`}`
`46`	`48`
`47`	`49`	`@Override`
`@@ -80,6 +82,6 @@ public ConnectorPageSource createPageSource(`
`80`	`82`	`return new EmptyPageSource();`
`81`	`83`	`}`
`82`	`84`
`83`		`- return new MongoPageSource(mongoSession, newTableHandle, handles.build());`
	`85`	`+ return new MongoPageSource(mongoSession, newTableHandle, handles.build(), implicitPrefix);`
`84`	`86`	`}`
`85`	`87`	`}`