Merge pull request #1260 from David-C-L/deltalake-scripting-updates

Deltalake scripting updates
RumbleDB · Sep 20, 2024 · 7ee88c3 · 7ee88c3
2 parents 41053cc + b02e732
commit 7ee88c3
Show file tree

Hide file tree

Showing 13 changed files with 155 additions and 30 deletions.
diff --git a/src/main/java/org/rumbledb/api/Rumble.java b/src/main/java/org/rumbledb/api/Rumble.java
@@ -52,11 +52,6 @@ public SequenceOfItems runQuery(String query) {
             this.configuration
         );
 
-        if (iterator.isUpdating()) {
-            PendingUpdateList pul = iterator.getPendingUpdateList(dynamicContext);
-            pul.applyUpdates(iterator.getMetadata());
-        }
-
         return new SequenceOfItems(iterator, dynamicContext, this.configuration);
     }
 
@@ -78,11 +73,6 @@ public SequenceOfItems runQuery(URI location) throws IOException {
             this.configuration
         );
 
-        if (iterator.isUpdating()) {
-            PendingUpdateList pul = iterator.getPendingUpdateList(dynamicContext);
-            pul.applyUpdates(iterator.getMetadata());
-        }
-
         return new SequenceOfItems(iterator, dynamicContext, this.configuration);
     }
 

diff --git a/src/main/java/org/rumbledb/api/SequenceOfItems.java b/src/main/java/org/rumbledb/api/SequenceOfItems.java
@@ -10,6 +10,7 @@
 import org.rumbledb.items.ItemFactory;
 import org.rumbledb.runtime.RuntimeIterator;
 
+import org.rumbledb.runtime.update.PendingUpdateList;
 import sparksoniq.spark.SparkSessionManager;
 
 /**
@@ -52,7 +53,7 @@ public SequenceOfItems(
      * Opens the iterator.
      */
     public void open() {
-        if (!this.isUpdating()) {
+        if (this.isMaterialisable()) {
             this.iterator.open(this.dynamicContext);
         }
         this.isOpen = true;
@@ -71,7 +72,7 @@ public boolean isOpen() {
      * Closes the iterator.
      */
     public void close() {
-        if (!this.isUpdating()) {
+        if (this.isOpen) {
             this.iterator.close();
         }
         this.isOpen = false;
@@ -83,7 +84,7 @@ public void close() {
      * @return true if there are more items, false otherwise.
      */
     public boolean hasNext() {
-        if (this.isUpdating()) {
+        if (!this.isMaterialisable()) {
             return false;
         }
         return this.iterator.hasNext();
@@ -96,7 +97,7 @@ public boolean hasNext() {
      * @return the next item.
      */
     public Item next() {
-        if (this.isUpdating()) {
+        if (!this.isMaterialisable()) {
             return ItemFactory.getInstance().createNullItem();
         }
         return this.iterator.next();
@@ -120,14 +121,32 @@ public boolean availableAsDataFrame() {
         return this.iterator.isDataFrame();
     }
 
+    /**
+     * Returns whether the iterator is updating
+     *
+     * @return true if updating; otherwise false.
+     */
+    public boolean availableAsPUL() {
+        return this.iterator.isUpdating();
+    }
+
+    /**
+     * Return whether the iterator of the sequence should be evaluated to materialise the sequence of items.
+     *
+     * @return true if materialisable; otherwise false
+     */
+    private boolean isMaterialisable() {
+        return !(this.availableAsPUL() && !this.iterator.isSequential());
+    }
+
     /**
      * Returns the sequence of items as an RDD of Items rather than iterating over them locally.
      * It is not possible to do so if the iterator is open.
      *
      * @return an RDD of Items.
      */
     public JavaRDD<Item> getAsRDD() {
-        if (this.isUpdating()) {
+        if (!this.isMaterialisable()) {
             return SparkSessionManager.getInstance().getJavaSparkContext().emptyRDD();
         }
         if (this.isOpen) {
@@ -143,7 +162,7 @@ public JavaRDD<Item> getAsRDD() {
      * @return a data frame.
      */
     public Dataset<Row> getAsDataFrame() {
-        if (this.isUpdating()) {
+        if (!this.isMaterialisable()) {
             return SparkSessionManager.getInstance().getOrCreateSession().emptyDataFrame();
         }
         if (this.isOpen) {
@@ -153,12 +172,11 @@ public Dataset<Row> getAsDataFrame() {
     }
 
     /**
-     * Returns whether the iterator is updating
-     *
-     * @return true if updating; otherwise false.
+     * Applies the PUL available when the iterator is updating.
      */
-    public boolean isUpdating() {
-        return this.iterator.isUpdating();
+    public void applyPUL() {
+        PendingUpdateList pul = this.iterator.getPendingUpdateList(this.dynamicContext);
+        pul.applyUpdates(this.iterator.getMetadata());
     }
 
     /*
@@ -168,7 +186,7 @@ public boolean isUpdating() {
      */
     public long populateList(List<Item> resultList) {
         resultList.clear();
-        if (this.isUpdating()) {
+        if (!this.isMaterialisable()) {
             return -1;
         }
         this.iterator.open(this.dynamicContext);
@@ -205,7 +223,7 @@ public long populateList(List<Item> resultList) {
 
     public long populateListWithWarningOnlyIfCapReached(List<Item> resultList) {
         if (this.availableAsRDD()) {
-            if (this.isUpdating()) {
+            if (!this.isMaterialisable()) {
                 return -1;
             }
             JavaRDD<Item> rdd = this.iterator.getRDD(this.dynamicContext);
@@ -215,5 +233,4 @@ public long populateListWithWarningOnlyIfCapReached(List<Item> resultList) {
         }
     }
 
-
 }
diff --git a/src/main/java/org/rumbledb/cli/JsoniqQueryExecutor.java b/src/main/java/org/rumbledb/cli/JsoniqQueryExecutor.java
@@ -34,6 +34,7 @@
 import org.rumbledb.optimizations.Profiler;
 import org.rumbledb.runtime.functions.input.FileSystemUtil;
 
+import org.rumbledb.runtime.update.PendingUpdateList;
 import sparksoniq.spark.SparkSessionManager;
 import java.io.IOException;
 import java.net.URI;
@@ -186,6 +187,10 @@ public List<Item> runQuery() throws IOException {
             }
         }
 
+        if (this.configuration.applyUpdates() && sequence.availableAsPUL() && outputPath != null) {
+            sequence.applyPUL();
+        }
+
         long endTime = System.currentTimeMillis();
         long totalTime = endTime - startTime;
         if (logPath != null) {
@@ -226,6 +231,9 @@ public long runInteractive(String query, List<Item> resultList) throws IOExcepti
         if (!sequence.availableAsRDD()) {
             return sequence.populateList(resultList);
         }
+        if (this.configuration.applyUpdates() && sequence.availableAsPUL()) {
+            sequence.applyPUL();
+        }
         resultList.clear();
         JavaRDD<Item> rdd = sequence.getAsRDD();
         return SparkSessionManager.collectRDDwithLimitWarningOnly(rdd, resultList);

diff --git a/src/main/java/org/rumbledb/compiler/ExpressionClassificationVisitor.java b/src/main/java/org/rumbledb/compiler/ExpressionClassificationVisitor.java
@@ -33,6 +33,7 @@
 import org.rumbledb.expressions.scripting.loops.ReturnStatementClause;
 import org.rumbledb.expressions.scripting.statement.StatementsAndExpr;
 import org.rumbledb.expressions.scripting.statement.StatementsAndOptionalExpr;
+import org.rumbledb.expressions.typing.TreatExpression;
 import org.rumbledb.expressions.update.AppendExpression;
 import org.rumbledb.expressions.update.CopyDeclaration;
 import org.rumbledb.expressions.update.DeleteExpression;
@@ -139,6 +140,18 @@ public ExpressionClassification visitCommaExpression(
         return expression.getExpressionClassification();
     }
 
+    // Region Typing
+
+    @Override
+    public ExpressionClassification visitTreatExpression(TreatExpression expression, ExpressionClassification argument) {
+        ExpressionClassification result = this.visit(expression.getMainExpression(), argument);
+        expression.setExpressionClassification(result);
+        return result;
+    }
+
+
+    // Endregion
+
     // Region FLWOR
 
     @Override

diff --git a/src/main/java/org/rumbledb/compiler/RuntimeIteratorVisitor.java b/src/main/java/org/rumbledb/compiler/RuntimeIteratorVisitor.java
@@ -1004,6 +1004,7 @@ public RuntimeIterator visitTreatExpression(TreatExpression expression, RuntimeI
         RuntimeIterator runtimeIterator = new TreatIterator(
                 childExpression,
                 expression.getSequenceType(),
+                expression.isUpdating(),
                 expression.errorCodeThatShouldBeThrown(),
                 expression.getStaticContextForRuntime(this.config, this.visitorConfig)
         );

diff --git a/src/main/java/org/rumbledb/config/RumbleRuntimeConfiguration.java b/src/main/java/org/rumbledb/config/RumbleRuntimeConfiguration.java
@@ -73,6 +73,7 @@ public class RumbleRuntimeConfiguration implements Serializable, KryoSerializabl
     private boolean nativeExecution;
     private boolean functionInlining;
     private boolean thirdFeature;
+    private boolean applyUpdates;
 
     private Map<String, String> shortcutMap;
     private Set<String> yesNoShortcuts;
@@ -422,6 +423,12 @@ public void init() {
             this.functionInlining = true;
         }
 
+        if (this.arguments.containsKey("apply-updates")) {
+            this.applyUpdates = this.arguments.get("apply-updates").equals("yes");
+        } else {
+            this.applyUpdates = false;
+        }
+
         if (this.arguments.containsKey("optimize-general-comparison-to-value-comparison")) {
             this.optimizeGeneralComparisonToValueComparison = this.arguments.get(
                 "optimize-general-comparison-to-value-comparison"
@@ -644,6 +651,14 @@ public void setFunctionInlining(boolean b) {
         this.functionInlining = b;
     }
 
+    public boolean applyUpdates() {
+        return this.applyUpdates;
+    }
+
+    public void setApplyUpdates(boolean b) {
+        this.applyUpdates = b;
+    }
+
     public boolean optimizeGeneralComparisonToValueComparison() {
         return this.optimizeGeneralComparisonToValueComparison;
     }

diff --git a/src/main/java/org/rumbledb/runtime/scripting/ProgramIterator.java b/src/main/java/org/rumbledb/runtime/scripting/ProgramIterator.java
@@ -105,6 +105,16 @@ protected Item nextLocal() {
         return null;
     }
 
+    @Override
+    public boolean isSequential() {
+        return this.statementsAndExprIterator.isSequential();
+    }
+
+    @Override
+    public boolean isUpdating() {
+        return this.statementsAndExprIterator.isUpdating();
+    }
+
     @Override
     public PendingUpdateList getPendingUpdateList(DynamicContext context) {
         if (!encounteredExitStatement) {

diff --git a/src/main/java/org/rumbledb/runtime/scripting/block/StatementsWithExprIterator.java b/src/main/java/org/rumbledb/runtime/scripting/block/StatementsWithExprIterator.java
@@ -156,6 +156,11 @@ public JSoundDataFrame getDataFrame(DynamicContext dynamicContext) {
         return exprIterator.getDataFrame(dynamicContext);
     }
 
+    @Override
+    public boolean isUpdating() {
+        return this.children.get(this.children.size() - 1).isUpdating();
+    }
+
     @Override
     public PendingUpdateList getPendingUpdateList(DynamicContext context) {
         RuntimeIterator exprIterator = this.children.get(this.children.size() - 1);

diff --git a/src/main/java/org/rumbledb/runtime/typing/TreatIterator.java b/src/main/java/org/rumbledb/runtime/typing/TreatIterator.java
@@ -23,6 +23,7 @@
 import org.rumbledb.runtime.HybridRuntimeIterator;
 import org.rumbledb.runtime.RuntimeIterator;
 import org.rumbledb.runtime.functions.sequences.general.TreatAsClosure;
+import org.rumbledb.runtime.update.PendingUpdateList;
 import org.rumbledb.types.ItemType;
 import org.rumbledb.types.ItemTypeFactory;
 import org.rumbledb.types.SequenceType;
@@ -51,12 +52,14 @@ public class TreatIterator extends HybridRuntimeIterator {
     public TreatIterator(
             RuntimeIterator iterator,
             SequenceType sequenceType,
+            boolean isUpdating,
             ErrorCode errorCode,
             RuntimeStaticContext staticContext
     ) {
         super(Collections.singletonList(iterator), staticContext);
         this.iterator = iterator;
         this.sequenceType = sequenceType;
+        this.isUpdating = isUpdating;
         this.errorCode = errorCode;
         if (!this.sequenceType.isEmptySequence()) {
             this.itemType = this.sequenceType.getItemType();
@@ -73,6 +76,15 @@ public TreatIterator(
         }
     }
 
+    public TreatIterator(
+            RuntimeIterator iterator,
+            SequenceType sequenceType,
+            ErrorCode errorCode,
+            RuntimeStaticContext staticContext
+    ) {
+        this(iterator, sequenceType, false, errorCode, staticContext);
+    }
+
     @Override
     public boolean hasNextLocal() {
         return this.hasNext;
@@ -226,6 +238,11 @@ public JSoundDataFrame getDataFrame(DynamicContext dynamicContext) {
         throw errorToThrow("" + dataItemType);
     }
 
+    @Override
+    public PendingUpdateList getPendingUpdateList(DynamicContext context) {
+        return iterator.getPendingUpdateList(context);
+    }
+
     /**
      * Converts a homogeneous RDD of atomic values to a DataFrame
      * 

diff --git a/src/test/java/iq/DeltaUpdateRuntimeTests.java b/src/test/java/iq/DeltaUpdateRuntimeTests.java
@@ -73,6 +73,8 @@ public RumbleRuntimeConfiguration getConfiguration() {
                     "yes",
                     "--show-error-info",
                     "yes",
+                    "--apply-updates",
+                    "yes",
                     "--materialization-cap",
                     "900000"
                 }
@@ -86,7 +88,9 @@ public RumbleRuntimeConfiguration getConfiguration() {
                 "--output-format",
                 "delta",
                 "--show-error-info",
-                "yes"
+                "yes",
+                "--apply-updates",
+                "yes",
             }
     );
 
@@ -97,7 +101,9 @@ public RumbleRuntimeConfiguration getConfiguration() {
                 "--output-format",
                 "delta",
                 "--show-error-info",
-                "yes"
+                "yes",
+                "--apply-updates",
+                "yes",
             }
     );
 
@@ -254,6 +260,7 @@ protected void checkExpectedOutput(
 
     protected String runIterators(SequenceOfItems sequence) {
         String actualOutput = getIteratorOutput(sequence);
+        applyPossibleUpdates(sequence);
         return actualOutput;
     }
 
@@ -304,10 +311,12 @@ && getConfiguration().getResultSizeCap() > 0)
 
     private String getRDDResults(SequenceOfItems sequence) {
         JavaRDD<Item> rdd = sequence.getAsRDD();
+        applyPossibleUpdates(sequence);
         JavaRDD<String> output = rdd.map(o -> o.serialize());
         List<String> collectedOutput = new ArrayList<String>();
         SparkSessionManager.collectRDDwithLimitWarningOnly(output, collectedOutput);
 
+
         if (collectedOutput.isEmpty()) {
             return "";
         }
@@ -329,4 +338,10 @@ private String getRDDResults(SequenceOfItems sequence) {
         return result;
     }
 
+    private void applyPossibleUpdates(SequenceOfItems sequence) {
+        if (getConfiguration().applyUpdates() && sequence.availableAsPUL()) {
+            sequence.applyPUL();
+        }
+    }
+
 }