fix grammar roundtrip for multi column (#2825)

finos · May 2, 2024 · 463f303 · 463f303
1 parent cad726c
commit 463f303
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/...java/org/finos/legend/engine/language/pure/grammar/from/domain/DomainParseTreeWalker.java b/...java/org/finos/legend/engine/language/pure/grammar/from/domain/DomainParseTreeWalker.java
@@ -1349,7 +1349,7 @@ else if (ctx.columnBuilders() != null)
             else if (colSpecArray != null)
             {
                 ColSpecArray colSpecArr = new ColSpecArray();
-                colSpecArr.colSpecs = ListIterate.collect(colSpecArray.oneColSpec(), c -> processOneColSpec(c, typeParametersNames, lambdaContext, space, wrapFlag, addLines, expressions));
+                colSpecArr.colSpecs = ListIterate.collect(colSpecArray.oneColSpec(), c -> processOneColSpec(c, typeParametersNames, lambdaContext, space, wrapFlag, addLines, Lists.mutable.of()));
                 result = DomainParseTreeWalker.wrapWithClassInstance(colSpecArr, walkerSourceInformation.getSourceInformation(colSpecArray), "colSpecArray");
                 ;
             }
@@ -1366,7 +1366,7 @@ else if (colSpecArray != null)
         return result;
     }
 
-    private ColSpec processOneColSpec(DomainParserGrammar.OneColSpecContext oneColSpec, List<String> typeParametersNames, LambdaContext lambdaContext, String space, boolean wrapFlag, boolean addLines, List<Variable> expressions)
+    private ColSpec processOneColSpec(DomainParserGrammar.OneColSpecContext oneColSpec, List<String> typeParametersNames, LambdaContext lambdaContext, String space, boolean wrapFlag, boolean addLines, List<Variable> params)
     {
         ValueSpecification result;
         ColSpec colSpec = new ColSpec();
@@ -1379,8 +1379,8 @@ private ColSpec processOneColSpec(DomainParserGrammar.OneColSpecContext oneColSp
         if (oneColSpec.lambdaParam() != null && oneColSpec.lambdaPipe() != null)
         {
             Variable var = this.lambdaParam(oneColSpec.lambdaParam().lambdaParamType(), oneColSpec.lambdaParam().identifier(), typeParametersNames, space);
-            expressions.add(var);
-            colSpec.function1 = this.lambdaPipe(oneColSpec.lambdaPipe(), oneColSpec.lambdaParam().getStart(), expressions, typeParametersNames, lambdaContext, space, wrapFlag, addLines);
+            params.add(var);
+            colSpec.function1 = this.lambdaPipe(oneColSpec.lambdaPipe(), oneColSpec.lambdaParam().getStart(), params, typeParametersNames, lambdaContext, space, wrapFlag, addLines);
             if (oneColSpec.extraFunction() != null)
             {
                 List<Variable> var2 = Lists.mutable.with(this.lambdaParam(oneColSpec.extraFunction().lambdaParam().lambdaParamType(), oneColSpec.extraFunction().lambdaParam().identifier(), typeParametersNames, space));

diff --git a/.../test/java/org/finos/legend/engine/language/pure/grammar/test/roundtrip/TestRelation.java b/.../test/java/org/finos/legend/engine/language/pure/grammar/test/roundtrip/TestRelation.java
@@ -68,4 +68,16 @@ public void testColumnArray()
         testLambda("|#>{path::Store.table}#->extend(~[a, b])");
     }
 
+    @Test
+    public void testSimpleProject()
+    {
+        testLambda("|test::Person.all()->project(~[mycol:x|$x.name])");
+    }
+
+    @Test
+    public void testProject()
+    {
+        testLambda("|test::Person.all()->project(~[first:x|$x.name, second:x|$x.val])");
+    }
+
 }