Fix keep/drop

vtl-engine/src/main/java/fr/insee/vtl/engine/visitors/ClauseVisitor.java

@@ -6,13 +6,15 @@
 
 import fr.insee.vtl.engine.VtlScriptEngine;
 import fr.insee.vtl.engine.exceptions.InvalidArgumentException;
+import fr.insee.vtl.engine.exceptions.UndefinedVariableException;
 import fr.insee.vtl.engine.exceptions.VtlRuntimeException;
 import fr.insee.vtl.engine.visitors.expression.ExpressionVisitor;
 import fr.insee.vtl.model.*;
 import fr.insee.vtl.model.exceptions.VtlScriptException;
 import fr.insee.vtl.parser.VtlBaseVisitor;
 import fr.insee.vtl.parser.VtlParser;
 import java.util.*;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.antlr.v4.runtime.ParserRuleContext;
 import org.antlr.v4.runtime.misc.Interval;
@@ -106,16 +108,59 @@ private static AggregationExpression convertToAggregation(
 
   @Override
   public DatasetExpression visitKeepOrDropClause(VtlParser.KeepOrDropClauseContext ctx) {
-    // Normalize to keep operation.
-    var keep = ctx.op.getType() == VtlParser.KEEP;
-    var names = ctx.componentID().stream().map(ClauseVisitor::getName).collect(Collectors.toSet());
-    List<String> columnNames =
-        datasetExpression.getDataStructure().values().stream()
-            .map(Dataset.Component::getName)
-            .filter(name -> keep == names.contains(name))
-            .collect(Collectors.toList());
 
-    return processingEngine.executeProject(datasetExpression, columnNames);
+    // The type of the op can either be KEEP or DROP.
+    boolean keep = ctx.op.getType() == VtlParser.KEEP;
+
+    // Dataset identifiers (role = IDENTIFIER)
+    Map<String, Dataset.Component> identifiers =
+        datasetExpression.getDataStructure().getIdentifiers().stream()
+            .collect(Collectors.toMap(Structured.Component::getName, Function.identity()));
+
+    var columns =
+        ctx.componentID().stream()
+            .collect(Collectors.toMap(ClauseVisitor::getName, Function.identity()));
+
+    var structure = datasetExpression.getDataStructure();
+
+    // Evaluate that all requested columns must exist in the dataset or raise an error
+    // TODO: Is that no handled already?
+    for (String col : columns.keySet()) {
+      if (!structure.containsKey(col)) {
+        throw new VtlRuntimeException(
+            new UndefinedVariableException(col, fromContext(columns.get(col))));
+      }
+    }
+
+    // VTL specification: identifiers must not appear explicitly in KEEP
+    // TODO: Use multi errors that noah created?
+    for (String col : columns.keySet()) {
+      if (structure.get(col).isIdentifier()) {
+        throw new VtlRuntimeException(
+            new InvalidArgumentException(
+                "cannot keep/drop identifiers", fromContext(columns.get(col))));
+      }
+    }
+
+    // Build result set:
+    //  + KEEP: identifiers + requested columns
+    //  + DROP: (all columns - requested) + identifiers
+    final Set<String> resultSet = new LinkedHashSet<>();
+    resultSet.addAll(identifiers.keySet());
+    if (keep) {
+      resultSet.addAll(columns.keySet());
+    } else {
+      for (String col : structure.keySet()) {
+        if (!columns.keySet().contains(col)) {
+          resultSet.add(col);
+        }
+      }
+    }
+
+    // Retrieve the output column names (identifiers + requested)
+    final List<String> outputColumns =
+        structure.keySet().stream().filter(resultSet::contains).collect(Collectors.toList());
+    return processingEngine.executeProject(datasetExpression, outputColumns);
   }
 
   @Override

vtl-engine/src/test/java/fr/insee/vtl/engine/VtlScriptEngineTest.java

@@ -33,24 +33,37 @@ public static <T extends Throwable> Condition<T> atPosition(
 
   public static <T extends Throwable> Condition<T> atPosition(
       Integer startLine, Integer endLine, Integer startColumn, Integer endColumn) {
-    return new Condition<>(
-        throwable -> {
-          var scriptException = (VtlScriptException) throwable;
-          var position = scriptException.getPosition();
-          return position.startLine().equals(startLine)
-              && position.endLine().equals(endLine)
-              && position.startColumn().equals(startColumn)
-              && position.endColumn().equals(endColumn);
-        },
-        "at position <%d:%d-%d:%d>",
-        startLine,
-        endLine,
-        startColumn,
-        endColumn);
-  }
-
-  public static <T extends Comparable<T>> Boolean test(T left, T right) {
-    return true;
+    return new Condition<T>() {
+      @Override
+      public boolean matches(T throwable) {
+        if (!(throwable instanceof VtlScriptException scriptException)) {
+          return false;
+        }
+        var position = scriptException.getPosition();
+        boolean matches =
+            position.startLine().equals(startLine)
+                && position.endLine().equals(endLine)
+                && position.startColumn().equals(startColumn)
+                && position.endColumn().equals(endColumn);
+
+        // Set description that includes actual position if it doesn't match
+        if (matches) {
+          describedAs("at position <%d:%d-%d:%d>", startLine, startColumn, endLine, endColumn);
+        } else {
+          describedAs(
+              "at position <%d:%d-%d:%d> but was <%d:%d-%d:%d>",
+              startLine,
+              startColumn,
+              endLine,
+              endColumn,
+              position.startLine(),
+              position.startColumn(),
+              position.endLine(),
+              position.endColumn());
+        }
+        return matches;
+      }
+    };
   }
 
   @BeforeEach

vtl-engine/src/test/java/fr/insee/vtl/engine/utils/dag/DagTest.java

@@ -199,11 +199,11 @@ void testDagSimpleExampleWithReordering() throws ScriptException {
   void testDagCycle() {
     final String script =
         """
-                            e := a;
-                            b := a;
-                            c := b;
-                            a := c;
-                            f := a;""";
+                        e := a;
+                        b := a;
+                        c := b;
+                        a := c;
+                        f := a;""";
 
     final Positioned.Position mainPosition = getPositionOfStatementInScript("a := c", script);
     final List<Positioned.Position> otherPositions =
@@ -231,14 +231,14 @@ void testDagCycle() {
   void testMultipleCycles() {
     final String script =
         """
-                            h := g;
-                            i := join(h, input_ds);
-                            g := i;
-                            e := a;
-                            b := a;
-                            c := b;
-                            a := c;
-                            f := a;""";
+                        h := g;
+                        i := join(h, input_ds);
+                        g := i;
+                        e := a;
+                        b := a;
+                        c := b;
+                        a := c;
+                        f := a;""";
 
     final Positioned.Position mainExceptionMainPosition =
         getPositionOfStatementInScript("g := i", script);
@@ -446,8 +446,8 @@ void testDagIfExpr() throws ScriptException {
     engine.getContext().setAttribute("ds_2", DatasetSamples.ds2, ScriptContext.ENGINE_SCOPE);
     engine.eval(
         "res := if ds1 > ds2 then ds1 else ds2; "
-            + "ds1 := ds_1[keep id, long1][rename long1 to bool_var]; "
-            + "ds2 := ds_2[keep id, long1][rename long1 to bool_var];");
+            + "ds1 := ds_1[keep long1][rename long1 to bool_var]; "
+            + "ds2 := ds_2[keep long1][rename long1 to bool_var];");
     var res = engine.getContext().getAttribute("res");
     assertThat(((Dataset) res).getDataAsMap())
         .containsExactlyInAnyOrder(
@@ -463,7 +463,7 @@ void testDagCaseExpr() throws ScriptException {
     engine.getContext().setAttribute("ds_2", DatasetSamples.ds2, ScriptContext.ENGINE_SCOPE);
     engine.eval(
         "res0 <- tmp0[calc c := case when long1 > 30 then \"ok\" else \"ko\"][drop long1]; "
-            + "tmp0 := ds_1[keep id, long1];");
+            + "tmp0 := ds_1[keep long1];");
     Object res0 = engine.getContext().getAttribute("res0");
     assertThat(((Dataset) res0).getDataAsMap())
         .containsExactlyInAnyOrder(
@@ -474,7 +474,7 @@ void testDagCaseExpr() throws ScriptException {
     assertThat(((Dataset) res0).getDataStructure().get("c").getType()).isEqualTo(String.class);
     engine.eval(
         "res1 <- tmp1[calc c := case when long1 > 30 then 1 else 0][drop long1]; "
-            + "tmp1 := ds_1[keep id, long1];");
+            + "tmp1 := ds_1[keep long1];");
     Object res1 = engine.getContext().getAttribute("res1");
     assertThat(((Dataset) res1).getDataAsMap())
         .containsExactlyInAnyOrder(
@@ -484,9 +484,9 @@ void testDagCaseExpr() throws ScriptException {
             Map.of("id", "Franck", "c", 1L));
     assertThat(((Dataset) res1).getDataStructure().get("c").getType()).isEqualTo(Long.class);
     engine.eval(
-        "tmp2_alt_ds1 := ds_1[keep id, long1][rename long1 to bool_var]; "
+        "tmp2_alt_ds1 := ds_1[keep long1][rename long1 to bool_var]; "
             + "res2 <- case when tmp2_alt_ds1 < 30 then tmp2_alt_ds1 else tmp2_alt_ds2; "
-            + "tmp2_alt_ds2 := ds_2[keep id, long1][rename long1 to bool_var];");
+            + "tmp2_alt_ds2 := ds_2[keep long1][rename long1 to bool_var];");
     Object resDs = engine.getContext().getAttribute("res2");
     assertThat(((Dataset) resDs).getDataAsMap())
         .containsExactlyInAnyOrder(
@@ -498,7 +498,7 @@ void testDagCaseExpr() throws ScriptException {
   @Test
   void testDagNvlExpr() throws ScriptException {
     engine.getContext().setAttribute("ds1", DatasetSamples.ds1, ScriptContext.ENGINE_SCOPE);
-    engine.eval("res <- nvl(tmp1[keep id, long1], 0); tmp1 := ds1;");
+    engine.eval("res <- nvl(tmp1[keep long1], 0); tmp1 := ds1;");
     var res = engine.getContext().getAttribute("res");
     assertThat(((Dataset) res).getDataAsMap())
         .containsExactlyInAnyOrder(
@@ -512,7 +512,7 @@ void testDagNvlExpr() throws ScriptException {
   @Test
   void testDagNvlImplicitCast() throws ScriptException {
     engine.getContext().setAttribute("ds1", DatasetSamples.ds1, ScriptContext.ENGINE_SCOPE);
-    engine.eval("res := nvl(tmp1[keep id, long1], 0.1); tmp1 <- ds1;");
+    engine.eval("res := nvl(tmp1[keep long1], 0.1); tmp1 <- ds1;");
     var res = engine.getContext().getAttribute("res");
     assertThat(((Dataset) res).getDataAsMap())
         .containsExactlyInAnyOrder(
@@ -528,7 +528,7 @@ void testDagUnaryExpr() throws ScriptException {
     ScriptContext context = engine.getContext();
 
     context.setAttribute("ds2", DatasetSamples.ds2, ScriptContext.ENGINE_SCOPE);
-    Object res = engine.eval("res := + tmp1[keep id, long1, double1]; tmp1 <- ds2;");
+    Object res = engine.eval("res := + tmp1[keep long1, double1]; tmp1 <- ds2;");
     assertThat(((Dataset) res).getDataAsMap())
         .containsExactlyInAnyOrder(
             Map.of("id", "Hadrien", "long1", 150L, "double1", 1.1D),
@@ -537,7 +537,7 @@ void testDagUnaryExpr() throws ScriptException {
     assertThat(((Dataset) res).getDataStructure().get("long1").getType()).isEqualTo(Long.class);
 
     context.setAttribute("ds2", DatasetSamples.ds2, ScriptContext.ENGINE_SCOPE);
-    Object res2 = engine.eval("res2 := - tmp2[keep id, long1, double1]; tmp2 := ds2;");
+    Object res2 = engine.eval("res2 := - tmp2[keep long1, double1]; tmp2 := ds2;");
     assertThat(((Dataset) res2).getDataAsMap())
         .containsExactlyInAnyOrder(
             Map.of("id", "Hadrien", "long1", -150L, "double1", -1.1D),

vtl-engine/src/test/java/fr/insee/vtl/engine/visitors/ClauseVisitorTest.java

@@ -87,7 +87,7 @@ public void testManyCalc() throws ScriptException {
   }
 
   @Test
-  public void testCalcRoleModifier() throws ScriptException {
+  public void testCalcRoleModifier_measuresAndAttributesOk() throws ScriptException {
     InMemoryDataset dataset =
         new InMemoryDataset(
             List.of(
@@ -194,7 +194,8 @@ public void testKeepDropClause() throws ScriptException {
     ScriptContext context = engine.getContext();
     context.setAttribute("ds1", dataset, ScriptContext.ENGINE_SCOPE);
 
-    engine.eval("ds2 := ds1[keep name, age];");
+    // KEEP: identifiers must not be listed explicitly; they are implicitly preserved.
+    engine.eval("ds2 := ds1[keep age];");
 
     assertThat(engine.getContext().getAttribute("ds2")).isInstanceOf(Dataset.class);
     assertThat(((Dataset) engine.getContext().getAttribute("ds2")).getDataAsMap())
@@ -213,6 +214,27 @@ public void testKeepDropClause() throws ScriptException {
             Map.of("name", "Franck", "age", 12L));
   }
 
+  /** KEEP/DROP: listing identifiers explicitly must raise a script error. */
+  @Test
+  public void testKeepDropClause_identifierExplicitShouldFail() {
+    InMemoryDataset dataset =
+        new InMemoryDataset(
+            List.of(
+                Map.of("name", "Hadrien", "age", 10L, "weight", 11L),
+                Map.of("name", "Nico", "age", 11L, "weight", 10L),
+                Map.of("name", "Franck", "age", 12L, "weight", 9L)),
+            Map.of("name", String.class, "age", Long.class, "weight", Long.class),
+            Map.of("name", Role.IDENTIFIER, "age", Role.MEASURE, "weight", Role.MEASURE));
+
+    ScriptContext context = engine.getContext();
+    context.setAttribute("ds1", dataset, ScriptContext.ENGINE_SCOPE);
+
+    assertThatThrownBy(() -> engine.eval("ds := ds1[keep name, age];"))
+        .isInstanceOf(VtlScriptException.class)
+        .hasMessage("cannot keep/drop identifiers")
+        .is(atPosition(0, 15, 19));
+  }
+
   @Test
   public void testAggregateType() {
     InMemoryDataset dataset =

vtl-engine/src/test/java/fr/insee/vtl/engine/visitors/expression/ArithmeticExprOrConcatTest.java

@@ -52,7 +52,7 @@ public void testPlus() throws ScriptException {
     assertThat(context.getAttribute("plus2")).isEqualTo(5.0);
 
     context.setAttribute("ds1", DatasetSamples.ds1, ScriptContext.ENGINE_SCOPE);
-    Object res = engine.eval("res := ds1[keep id, long1, long2] + ds1[keep id, long1, long2];");
+    Object res = engine.eval("res := ds1[keep long1, long2] + ds1[keep long1, long2];");
     assertThat(((Dataset) res).getDataAsMap())
         .containsExactlyInAnyOrder(
             Map.of("id", "Toto", "long1", 60L, "long2", 600L),
@@ -75,7 +75,7 @@ public void testMinus() throws ScriptException {
 
     context.setAttribute("ds1", DatasetSamples.ds1, ScriptContext.ENGINE_SCOPE);
     context.setAttribute("ds2", DatasetSamples.ds2, ScriptContext.ENGINE_SCOPE);
-    Object res = engine.eval("res := ds2[keep id, long1] - ds1[keep id, long1] + 1;");
+    Object res = engine.eval("res := ds2[keep long1] - ds1[keep long1] + 1;");
     assertThat(((Dataset) res).getDataAsMap())
         .containsExactlyInAnyOrder(
             Map.of("id", "Hadrien", "long1", 141L),
@@ -92,7 +92,7 @@ public void testConcat() throws ScriptException {
 
     context.setAttribute("ds1", DatasetSamples.ds1, ScriptContext.ENGINE_SCOPE);
     context.setAttribute("ds2", DatasetSamples.ds2, ScriptContext.ENGINE_SCOPE);
-    Object res = engine.eval("res := ds2[keep id, string1] || \" \" || ds1[keep id, string1];");
+    Object res = engine.eval("res := ds2[keep string1] || \" \" || ds1[keep string1];");
     assertThat(((Dataset) res).getDataAsMap())
         .containsExactlyInAnyOrder(
             Map.of("id", "Hadrien", "string1", "hadrien hadrien"),

vtl-engine/src/test/java/fr/insee/vtl/engine/visitors/expression/ArithmeticExprTest.java

@@ -64,8 +64,7 @@ public void testArithmeticExpr() throws ScriptException {
 
     context.setAttribute("ds1", DatasetSamples.ds1, ScriptContext.ENGINE_SCOPE);
     context.setAttribute("ds2", DatasetSamples.ds2, ScriptContext.ENGINE_SCOPE);
-    Object res =
-        engine.eval("res := round(ds1[keep id, long1, double1] * ds2[keep id, long1, double1]);");
+    Object res = engine.eval("res := round(ds1[keep long1, double1] * ds2[keep long1, double1]);");
     assertThat(((Dataset) res).getDataAsMap())
         .containsExactlyInAnyOrder(
             Map.of("id", "Hadrien", "long1", 1500.0, "double1", 1.0),
@@ -83,8 +82,7 @@ public void testArithmeticExpr() throws ScriptException {
     engine.eval("div4 := 3.0 / 1.5;");
     assertThat(context.getAttribute("div4")).isEqualTo(2.0);
 
-    res =
-        engine.eval("res2 := round(ds1[keep id, long1, double1] / ds2[keep id, long1, double1]);");
+    res = engine.eval("res2 := round(ds1[keep long1, double1] / ds2[keep long1, double1]);");
     assertThat(((Dataset) res).getDataAsMap())
         .containsExactlyInAnyOrder(
             Map.of("id", "Hadrien", "long1", 0.0, "double1", 1.0),

vtl-engine/src/test/java/fr/insee/vtl/engine/visitors/expression/BooleanExprTest.java

@@ -100,8 +100,8 @@ public void testOnDatasets() throws ScriptException {
     context.setAttribute("ds_1", DatasetSamples.ds1, ScriptContext.ENGINE_SCOPE);
     context.setAttribute("ds_2", DatasetSamples.ds1, ScriptContext.ENGINE_SCOPE);
     engine.eval(
-        "ds1 := ds_1[keep id, bool2][rename bool2 to bool1]; "
-            + "ds2 := ds_2[keep id, bool1]; "
+        "ds1 := ds_1[keep bool2][rename bool2 to bool1]; "
+            + "ds2 := ds_2[keep  bool1]; "
             + "andDs := ds1 and ds2; "
             + "orDs := ds1 or ds2; "
             + "xorDs := ds1 xor ds2; ");