Add Utest about dividing zero (#3128)

* Add Utest about dividing zero * add Utest and zero check of LocalData * add Utest and zero check of LocalData * change * Add Utest about dividing zero * fix test
dding3 · Jun 2, 2021 · 007769c · 007769c
1 parent 83d901b
commit 007769c
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/scala/dllib/src/main/scala/com/intel/analytics/bigdl/dllib/feature/dataset/DataSet.scala b/scala/dllib/src/main/scala/com/intel/analytics/bigdl/dllib/feature/dataset/DataSet.scala
@@ -281,6 +281,9 @@ class CachedDistriDataSet[T: ClassTag] private[dataset]
         override def next(): T = {
           val i = _offset.getAndIncrement()
           if (_train) {
+            require(localData.length != 0,
+              "dataset on this executor is empty, please increase " +
+                "your dataset size or decrease your number of executors.")
             localData(indexes(i % localData.length))
           } else {
             if (i < localData.length) {

diff --git a/scala/dllib/src/test/scala/com/intel/analytics/bigdl/dllib/dataset/DataSetSpec.scala b/scala/dllib/src/test/scala/com/intel/analytics/bigdl/dllib/dataset/DataSetSpec.scala
@@ -470,4 +470,21 @@ class DataSetSpec extends SparkContextLifeCycle with Matchers {
     img.content((100 + 100 * 213) * 3 + 1) should be(30 / 255f)
     img.content((100 + 100 * 213) * 3) should be(36 / 255f)
   }
+
+  "small dataset" should "try catch" in {
+    val errormsg = "dataset on this executor is empty, please increase " +
+      "your dataset size or decrease your number of executors."
+    try {
+      val data = Array.empty[Int]
+      val dataset = DataSet.rdd(sc.parallelize(data, 3)).data(train = true)
+      dataset.mapPartitions(iter => {
+        Iterator.single(iter.next())
+      }).collect()(0)
+    }
+    catch {
+      case x: org.apache.spark.SparkException =>
+        x.getMessage().contains(errormsg) should be(true)
+    }
+  }
 }
+