record-related functions enhancements

PastorGL · Jan 15, 2025 · 6943aa4 · 6943aa4
1 parent 2d03f95
commit 6943aa4
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 13 deletions.
diff --git a/cli/src/main/java/io/github/pastorgl/datacooker/cli/repl/REPL.java b/cli/src/main/java/io/github/pastorgl/datacooker/cli/repl/REPL.java
@@ -355,6 +355,10 @@ public void loop() throws Exception {
                                             sb.append("\tRecord Key (implicit). Additional in description\n");
                                             break;
                                         }
+                                        case Function.RECORD_LEVEL: {
+                                            sb.append("\tRecord level with any number of arguments\n");
+                                            break;
+                                        }
                                         case Function.RECORD_OBJECT: {
                                             sb.append("\tRecord Object: " + ei.argTypes[0] + " (implicit). Additional in description\n");
                                             break;

diff --git a/cli/src/test/resources/function.vm b/cli/src/test/resources/function.vm
@@ -11,13 +11,15 @@
 #elseif( $op.arity == 0 )
 No arguments
 #elseif( $op.arity == -1 )
-Any number of arguments: <code>${op.argTypes[0]}</code>
+Any number of arguments of type <code>${op.argTypes[0]}</code>
 #elseif( $op.arity == -2 )
 Record Key (implicit). See description for any additional arguments
 #elseif( $op.arity == -3 )
 Record Object of type <code>${op.argTypes[0]}</code> (implicit). See description for any additional arguments
 #elseif( $op.arity == -4 )
 Record Key (implicit), Record Object (implicit). See description for any additional arguments
+#elseif( $op.arity == -5 )
+Record level with any number of arguments. See description
 #end
 </p>
 

diff --git a/...s/src/main/java/io/github/pastorgl/datacooker/commons/functions/RecordLevelFunctions.java b/...s/src/main/java/io/github/pastorgl/datacooker/commons/functions/RecordLevelFunctions.java
@@ -6,8 +6,8 @@
 
 import io.github.pastorgl.datacooker.data.DataRecord;
 import io.github.pastorgl.datacooker.scripting.Evaluator;
-import io.github.pastorgl.datacooker.scripting.Function;
 import io.github.pastorgl.datacooker.scripting.Function.RecordKey;
+import io.github.pastorgl.datacooker.scripting.Function.RecordLevel;
 import io.github.pastorgl.datacooker.scripting.Function.RecordObject;
 
 import java.util.Deque;
@@ -83,12 +83,9 @@ public String descr() {
         }
     }
 
-    public static class PIVOT extends Function.WholeRecord<Object[]> {
+    public static class PIVOT extends RecordLevel<Object[]> {
         @Override
         public Object[] call(Deque<Object> args) {
-            //discard bot record and key; this should work only in SELECT context for whole records
-            args.pop();
-            args.pop();
             return Evaluator.popArray(args).data();
         }
 

diff --git a/commons/src/main/java/io/github/pastorgl/datacooker/commons/functions/StringFunctions.java b/commons/src/main/java/io/github/pastorgl/datacooker/commons/functions/StringFunctions.java
@@ -184,7 +184,7 @@ public String descr() {
         }
     }
 
-    public static class Format extends ArbitrAry<String, String> {
+    public static class Format extends ArbitrAry<String, Object> {
         @Override
         public String call(Deque<Object> args) {
             String format = Evaluator.popString(args);

diff --git a/commons/src/main/java/io/github/pastorgl/datacooker/data/DataContext.java b/commons/src/main/java/io/github/pastorgl/datacooker/data/DataContext.java
@@ -236,7 +236,7 @@ public StreamInfo alterDataStream(String dsName, StreamConverter converter, Map<
                 }
 
                 return new DataStreamBuilder(dsName, ds.attributes())
-                        .altered("KEY", ds)
+                        .altered("KEY" + (shuffle ? " PARTITION" : ""), ds)
                         .keyExpr(ke)
                         .build(reKeyed);
             };

diff --git a/commons/src/main/java/io/github/pastorgl/datacooker/scripting/Function.java b/commons/src/main/java/io/github/pastorgl/datacooker/scripting/Function.java
@@ -7,19 +7,29 @@
 import io.github.pastorgl.datacooker.data.DataRecord;
 
 public abstract class Function<R> implements Evaluator<R> {
+    // record-related arities, allowed only in query context
+    public static final int RECORD_LEVEL = -5;
     public static final int WHOLE_RECORD = -4;
     public static final int RECORD_OBJECT = -3;
     public static final int RECORD_KEY = -2;
+    // anything above 'arbitrary' is allowed everywhere
     public static final int ARBITR_ARY = -1;
     public static final int NO_ARGS = 0;
 
     public Function() {
     }
 
-    public static abstract class RecordKey<R> extends Function<R> {
+    public static abstract class RecordLevel<R> extends Function<R> {
         @Override
         public int arity() {
-            return RECORD_KEY;
+            return RECORD_LEVEL;
+        }
+    }
+
+    public static abstract class WholeRecord<R, REC extends DataRecord<?>> extends Function<R> {
+        @Override
+        public int arity() {
+            return WHOLE_RECORD;
         }
     }
 
@@ -30,10 +40,10 @@ public int arity() {
         }
     }
 
-    public static abstract class WholeRecord<R> extends Function<R> {
+    public static abstract class RecordKey<R> extends Function<R> {
         @Override
         public int arity() {
-            return WHOLE_RECORD;
+            return RECORD_KEY;
         }
     }
 

diff --git a/commons/src/main/java/io/github/pastorgl/datacooker/scripting/TDL4Interpreter.java b/commons/src/main/java/io/github/pastorgl/datacooker/scripting/TDL4Interpreter.java
@@ -819,7 +819,12 @@ private List<Expressions.ExprItem<?>> expression(List<ParseTree> exprChildren, E
                     throw new RuntimeException("Unknown function token " + exprItem.getText());
                 } else {
                     int arity = ef.arity();
-                    if (arity == Function.ARBITR_ARY) {
+
+                    if ((arity < Function.ARBITR_ARY) && (rules != ExpressionRules.QUERY)) {
+                        throw new RuntimeException("Record-related function " + ef.name() + " can't be called outside of query context");
+                    }
+
+                    if ((arity == Function.ARBITR_ARY) || (arity == Function.RECORD_LEVEL)) {
                         items.add(Expressions.stackGetter(funcCall.expression().size()));
                     } else if (arity > 0) {
                         items.add(Expressions.stackGetter(arity));
@@ -837,6 +842,11 @@ private List<Expressions.ExprItem<?>> expression(List<ParseTree> exprChildren, E
                     throw new RuntimeException("Unknown function token " + exprItem.getText());
                 } else {
                     int arity = ef.arity();
+
+                    if ((arity < Function.ARBITR_ARY) && (rules != ExpressionRules.QUERY)) {
+                        throw new RuntimeException("Record-related function " + ef.name() + " can't be called outside of query context");
+                    }
+
                     switch (arity) {
                         case Function.RECORD_KEY: {
                             items.add(Expressions.keyItem(funcAttr.attr_expr().size()));
@@ -850,6 +860,7 @@ private List<Expressions.ExprItem<?>> expression(List<ParseTree> exprChildren, E
                             items.add(Expressions.recItem(funcAttr.attr_expr().size()));
                             break;
                         }
+                        case Function.RECORD_LEVEL:
                         case Function.ARBITR_ARY: {
                             items.add(Expressions.stackGetter(funcAttr.attr_expr().size()));
                             break;