diff --git a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/DataMapManager.java b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/DataMapManager.java
index d19a0a9ef9..6a07dd5e87 100644
--- a/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/DataMapManager.java
+++ b/flow-model-generator/modules/flow-model-generator-core/src/main/java/io/ballerina/flowmodelgenerator/core/DataMapManager.java
@@ -444,7 +444,7 @@ private TargetNode getTargetNode(Node parentNode, String targetField, SemanticMo
         String[] fieldSplits = targetField.split(DOT);
         int idx = 1;
         if (initializer.kind() == SyntaxKind.QUERY_EXPRESSION) {
-            if (fieldSplits.length == 2 && fieldSplits[1].equals("0")) {
+            if (fieldSplits.length >= 2 && fieldSplits[1].equals(ZERO)) {
                 idx = 2;
             }
         }
diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/java/io/ballerina/flowmodelgenerator/extension/DataMappingModelTest.java b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/java/io/ballerina/flowmodelgenerator/extension/DataMappingModelTest.java
index 6054742e1e..11020bf65b 100644
--- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/java/io/ballerina/flowmodelgenerator/extension/DataMappingModelTest.java
+++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/java/io/ballerina/flowmodelgenerator/extension/DataMappingModelTest.java
@@ -138,6 +138,8 @@ protected Object[] getConfigsList() {
                 {Path.of("variable50.json")},
                 {Path.of("variable51.json")},
                 {Path.of("variable52.json")},
+                {Path.of("query18.json")},
+                {Path.of("query19.json")},
         };
     }
 
diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query13.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query13.json
index 489c7576e6..d4aec594e9 100644
--- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query13.json
+++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query13.json
@@ -21,26 +21,9 @@
     "offset": 4
   },
   "propertyKey": "expression",
-  "targetField": "myTypes4.0.0",
+  "targetField": "myTypes4.0",
   "model": {
     "inputs": [
-      {
-        "member": {
-          "fields": [],
-          "name": "myType3",
-          "displayName": "<myType3Item>",
-          "typeName": "MyType",
-          "kind": "record",
-          "ref": "2079019383"
-        },
-        "name": "myType3",
-        "displayName": "myType3",
-        "typeName": "MyType[]",
-        "kind": "array",
-        "category": "variable",
-        "focusExpression": "myTypes3",
-        "ref": "2079019383"
-      },
       {
         "member": {
           "fields": [],
@@ -106,54 +89,75 @@
         "ref": "-1223228350"
       },
       {
-        "fields": [],
-        "name": "y",
-        "displayName": "y",
-        "typeName": "MyType",
-        "kind": "record",
-        "focusExpression": "myType3",
+        "member": {
+          "fields": [],
+          "name": "myType3",
+          "displayName": "<myType3Item>",
+          "typeName": "MyType",
+          "kind": "record",
+          "ref": "2079019383"
+        },
+        "name": "myType3",
+        "displayName": "myType3",
+        "typeName": "MyType[]",
+        "kind": "array",
+        "focusExpression": "myTypes3",
         "ref": "2079019383"
       }
     ],
     "output": {
-      "fields": [],
+      "member": {
+        "member": {
+          "fields": [],
+          "name": "0",
+          "displayName": "<0ItemItem>",
+          "typeName": "MyType",
+          "kind": "record",
+          "ref": "2079019383"
+        },
+        "name": "0",
+        "displayName": "<0Item>",
+        "typeName": "MyType[]",
+        "kind": "array",
+        "ref": "2079019383"
+      },
       "name": "0",
       "displayName": "0",
-      "typeName": "record {|string name1; int i;|}",
-      "kind": "record",
-      "ref": "-353186209"
+      "typeName": "MyType[][]",
+      "kind": "array",
+      "ref": "2079019383"
     },
     "mappings": [
       {
         "output": "0",
         "inputs": [
-          "y"
+          "myType3"
         ],
-        "expression": "y",
+        "expression": "from var y in myType3 select y",
         "diagnostics": [],
         "elements": [],
-        "isQueryExpression": false,
+        "isQueryExpression": true,
         "isFunctionCall": false
       }
     ],
     "query": {
       "output": "0",
       "inputs": [
-        "myType3"
+        "myTypes3"
       ],
       "fromClause": {
         "type": "from",
         "properties": {
-          "name": "y",
-          "type": "MyType",
-          "expression": "myType3"
+          "name": "myType3",
+          "type": "MyType[]",
+          "expression": "myTypes3"
         }
       },
       "intermediateClauses": [],
       "resultClause": {
         "type": "select",
         "properties": {
-          "expression": "y"
+          "expression": "from var y in myType3 select y"
         }
       }
     },
@@ -206,26 +210,6 @@
         ],
         "typeName": "Types",
         "kind": "record"
-      },
-      "-353186209": {
-        "fields": [
-          {
-            "name": "name1",
-            "displayName": "name1",
-            "typeName": "string",
-            "kind": "string",
-            "optional": false
-          },
-          {
-            "name": "i",
-            "displayName": "i",
-            "typeName": "int",
-            "kind": "int",
-            "optional": false
-          }
-        ],
-        "typeName": "record {|string name1; int i;|}",
-        "kind": "record"
       }
     }
   }
diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query15.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query15.json
index 489c7576e6..f0a3cf2c90 100644
--- a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query15.json
+++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query15.json
@@ -4,7 +4,7 @@
   "codedata": {
     "node": "VARIABLE",
     "lineRange": {
-      "fileName": "query7.bal",
+      "fileName": "query8.bal",
       "startLine": {
         "line": 14,
         "offset": 4
@@ -116,12 +116,19 @@
       }
     ],
     "output": {
-      "fields": [],
+      "member": {
+        "fields": [],
+        "name": "0",
+        "displayName": "<0Item>",
+        "typeName": "MyType",
+        "kind": "record",
+        "ref": "2079019383"
+      },
       "name": "0",
       "displayName": "0",
-      "typeName": "record {|string name1; int i;|}",
-      "kind": "record",
-      "ref": "-353186209"
+      "typeName": "MyType[]",
+      "kind": "array",
+      "ref": "2079019383"
     },
     "mappings": [
       {
@@ -206,26 +213,6 @@
         ],
         "typeName": "Types",
         "kind": "record"
-      },
-      "-353186209": {
-        "fields": [
-          {
-            "name": "name1",
-            "displayName": "name1",
-            "typeName": "string",
-            "kind": "string",
-            "optional": false
-          },
-          {
-            "name": "i",
-            "displayName": "i",
-            "typeName": "int",
-            "kind": "int",
-            "optional": false
-          }
-        ],
-        "typeName": "record {|string name1; int i;|}",
-        "kind": "record"
       }
     }
   }
diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query18.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query18.json
new file mode 100644
index 0000000000..2ae6cb6c8c
--- /dev/null
+++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query18.json
@@ -0,0 +1,140 @@
+{
+  "source": "query9.bal",
+  "description": "Sample diagram node",
+  "codedata": {
+    "node": "VARIABLE",
+    "lineRange": {
+      "fileName": "query9.bal",
+      "startLine": {
+        "line": 6,
+        "offset": 4
+      },
+      "endLine": {
+        "line": 8,
+        "offset": 42
+      }
+    },
+    "sourceCode": ""
+  },
+  "position": {
+    "line": 6,
+    "offset": 4
+  },
+  "propertyKey": "expression",
+  "targetField": "arr4.0.0",
+  "model": {
+    "inputs": [
+      {
+        "member": {
+          "member": {
+            "fields": [],
+            "name": "arr3",
+            "displayName": "<arr3ItemItem>",
+            "typeName": "RecType",
+            "kind": "record",
+            "ref": "-1897494121"
+          },
+          "name": "arr3",
+          "displayName": "<arr3Item>",
+          "typeName": "RecType[]",
+          "kind": "array",
+          "ref": "-1897494121"
+        },
+        "name": "arr3",
+        "displayName": "arr3",
+        "typeName": "RecType[][]",
+        "kind": "array",
+        "category": "variable",
+        "ref": "-1897494121"
+      },
+      {
+        "member": {
+          "fields": [],
+          "name": "arr3Item",
+          "displayName": "<arr3ItemItem>",
+          "typeName": "RecType",
+          "kind": "record",
+          "ref": "-1897494121"
+        },
+        "name": "arr3Item",
+        "displayName": "arr3Item",
+        "typeName": "RecType[]",
+        "kind": "array",
+        "category": "variable",
+        "focusExpression": "arr3",
+        "ref": "-1897494121"
+      },
+      {
+        "fields": [],
+        "name": "arr3ItemItem",
+        "displayName": "arr3ItemItem",
+        "typeName": "RecType",
+        "kind": "record",
+        "focusExpression": "arr3Item",
+        "ref": "-1897494121"
+      }
+    ],
+    "output": {
+      "member": {
+        "fields": [],
+        "name": "0",
+        "displayName": "<0Item>",
+        "typeName": "RecType",
+        "kind": "record",
+        "ref": "-1897494121"
+      },
+      "name": "0",
+      "displayName": "0",
+      "typeName": "RecType[]",
+      "kind": "array",
+      "ref": "-1897494121"
+    },
+    "mappings": [
+      {
+        "output": "0.i",
+        "inputs": [],
+        "expression": "0",
+        "diagnostics": [],
+        "elements": [],
+        "isQueryExpression": false,
+        "isFunctionCall": false
+      }
+    ],
+    "query": {
+      "output": "0",
+      "inputs": [
+        "arr3Item"
+      ],
+      "fromClause": {
+        "type": "from",
+        "properties": {
+          "name": "arr3ItemItem",
+          "type": "RecType",
+          "expression": "arr3Item"
+        }
+      },
+      "intermediateClauses": [],
+      "resultClause": {
+        "type": "select",
+        "properties": {
+          "expression": "{i: 0}"
+        }
+      }
+    },
+    "refs": {
+      "-1897494121": {
+        "fields": [
+          {
+            "name": "i",
+            "displayName": "i",
+            "typeName": "int",
+            "kind": "int",
+            "optional": false
+          }
+        ],
+        "typeName": "RecType",
+        "kind": "record"
+      }
+    }
+  }
+}
diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query19.json b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query19.json
new file mode 100644
index 0000000000..3a9ae4ba75
--- /dev/null
+++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/config/query19.json
@@ -0,0 +1,163 @@
+{
+  "source": "query10.bal",
+  "description": "Sample diagram node",
+  "codedata": {
+    "node": "VARIABLE",
+    "lineRange": {
+      "fileName": "query10.bal",
+      "startLine": {
+        "line": 13,
+        "offset": 4
+      },
+      "endLine": {
+        "line": 15,
+        "offset": 63
+      }
+    },
+    "sourceCode": ""
+  },
+  "position": {
+    "line": 13,
+    "offset": 4
+  },
+  "propertyKey": "expression",
+  "targetField": "myRecType.recType.0",
+  "model": {
+    "inputs": [
+      {
+        "member": {
+          "member": {
+            "fields": [],
+            "name": "arr3",
+            "displayName": "<arr3ItemItem>",
+            "typeName": "RecType",
+            "kind": "record",
+            "ref": "-1897494121"
+          },
+          "name": "arr3",
+          "displayName": "<arr3Item>",
+          "typeName": "RecType[]",
+          "kind": "array",
+          "ref": "-1897494121"
+        },
+        "name": "arr3",
+        "displayName": "arr3",
+        "typeName": "RecType[][]",
+        "kind": "array",
+        "category": "variable",
+        "ref": "-1897494121"
+      },
+      {
+        "member": {
+          "fields": [],
+          "name": "arr3Item",
+          "displayName": "<arr3ItemItem>",
+          "typeName": "RecType",
+          "kind": "record",
+          "ref": "-1897494121"
+        },
+        "name": "arr3Item",
+        "displayName": "arr3Item",
+        "typeName": "RecType[]",
+        "kind": "array",
+        "category": "variable",
+        "focusExpression": "arr3",
+        "ref": "-1897494121"
+      },
+      {
+        "member": {
+          "member": {
+            "fields": [],
+            "name": "arr4",
+            "displayName": "<arr4ItemItem>",
+            "typeName": "RecType",
+            "kind": "record",
+            "ref": "-1897494121"
+          },
+          "name": "arr4",
+          "displayName": "<arr4Item>",
+          "typeName": "RecType[]",
+          "kind": "array",
+          "ref": "-1897494121"
+        },
+        "name": "arr4",
+        "displayName": "arr4",
+        "typeName": "RecType[][]",
+        "kind": "array",
+        "category": "variable",
+        "ref": "-1897494121"
+      },
+      {
+        "fields": [],
+        "name": "arr3ItemItem",
+        "displayName": "arr3ItemItem",
+        "typeName": "RecType",
+        "kind": "record",
+        "focusExpression": "arr3Item",
+        "ref": "-1897494121"
+      }
+    ],
+    "output": {
+      "member": {
+        "fields": [],
+        "name": "0",
+        "displayName": "<0Item>",
+        "typeName": "RecType",
+        "kind": "record",
+        "ref": "-1897494121"
+      },
+      "name": "0",
+      "displayName": "0",
+      "typeName": "RecType[]",
+      "kind": "array",
+      "ref": "-1897494121"
+    },
+    "mappings": [
+      {
+        "output": "0.i",
+        "inputs": [],
+        "expression": "0",
+        "diagnostics": [],
+        "elements": [],
+        "isQueryExpression": false,
+        "isFunctionCall": false
+      }
+    ],
+    "query": {
+      "output": "0",
+      "inputs": [
+        "arr3Item"
+      ],
+      "fromClause": {
+        "type": "from",
+        "properties": {
+          "name": "arr3ItemItem",
+          "type": "RecType",
+          "expression": "arr3Item"
+        }
+      },
+      "intermediateClauses": [],
+      "resultClause": {
+        "type": "select",
+        "properties": {
+          "expression": "{i: 0}"
+        }
+      }
+    },
+    "refs": {
+      "-1897494121": {
+        "fields": [
+          {
+            "name": "i",
+            "displayName": "i",
+            "typeName": "int",
+            "kind": "int",
+            "optional": false
+          }
+        ],
+        "typeName": "RecType",
+        "kind": "record"
+      }
+    }
+  }
+}
diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/source/query10.bal b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/source/query10.bal
new file mode 100644
index 0000000000..1cfdf8c4c5
--- /dev/null
+++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/source/query10.bal
@@ -0,0 +1,17 @@
+type RecType record {|
+    int i;
+|};
+
+type MyRecType record {|
+    RecType[][] recType;
+|};
+
+public function main() {
+    RecType[][] arr3 = [];
+    RecType[][] arr4 = from var arr3Item in arr3
+                        select from var arr3ItemItem in arr3Item
+                            select {i: 0};
+    MyRecType myRecType = {recType: from var arr3Item in arr3
+                                        select from var arr3ItemItem in arr3Item
+                                                select {i: 0}};
+}
diff --git a/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/source/query9.bal b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/source/query9.bal
new file mode 100644
index 0000000000..98a865a92e
--- /dev/null
+++ b/flow-model-generator/modules/flow-model-generator-ls-extension/src/test/resources/data_mapper_model/source/query9.bal
@@ -0,0 +1,10 @@
+type RecType record {|
+    int i;
+|};
+
+public function main() {
+    RecType[][] arr3 = [];
+    RecType[][] arr4 = from var arr3Item in arr3
+                        select from var arr3ItemItem in arr3Item
+                            select {i: 0};
+}