[SW-1623] Fix tests causing intermittent NPE in PySparkling with roll…

…ups on external backend (#1534) (cherry picked from commit f7e5031)
h2oai · Sep 18, 2019 · dce70c0 · dce70c0
1 parent 0700c45
commit dce70c0
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/examples/scripts/hamOrSpam.script.scala b/examples/scripts/hamOrSpam.script.scala
@@ -141,7 +141,9 @@ def isSpam(msg: String,
     hashingTF.transform (
       tokenize (msgRdd))).map(v => ("?", v)).toDF("target", "fv")
   val msgTable: H2OFrame = h2oContext.asH2OFrame(msgVector)
-  msgTable.remove(0) // remove first column
+  H2OFrameSupport.withLockAndUpdate(msgTable) {
+    _.remove(0) // remove first column
+  }
   val prediction = dlModel.score(msgTable)
   //println(prediction)
   prediction.vecs()(1).at(0) < hamThreshold

diff --git a/examples/src/main/scala/org/apache/spark/examples/h2o/HamOrSpamDemo.scala b/examples/src/main/scala/org/apache/spark/examples/h2o/HamOrSpamDemo.scala
@@ -180,7 +180,9 @@ object HamOrSpamDemo extends SparkContextSupport with ModelMetricsSupport with H
       hashingTF.transform(
         tokenize(msgRdd))).map(v => SMS("?", v)).toDF
     val msgTable: H2OFrame = msgVector
-    msgTable.remove(0) // remove first column
+    H2OFrameSupport.withLockAndUpdate(msgTable) {
+      _.remove(0) // remove first column
+    }
     val prediction = dlModel.score(msgTable)
     //println(prediction)
     prediction.vecs()(1).at(0) < hamThreshold