use java map (#21)

* use java map * sync with main * remove unused imports * update filter selection * fix exception message
teragrep · Dec 7, 2023 · 9e62f53 · 9e62f53
1 parent 8a38613
commit 9e62f53
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 12 deletions.
diff --git a/src/main/scala/com/teragrep/functions/dpf_03/BloomFilterAggregator.scala b/src/main/scala/com/teragrep/functions/dpf_03/BloomFilterAggregator.scala
@@ -55,8 +55,9 @@ import org.apache.spark.util.sketch.BloomFilter
 
 import scala.collection.mutable
 import scala.reflect.ClassTag
+import scala.collection.JavaConverters._
 
-class BloomFilterAggregator(final val columnName: String, final val estimateName: String, sizeMap: mutable.SortedMap[Long, Double]) extends Aggregator[Row, BloomFilter, Array[Byte]]
+class BloomFilterAggregator(final val columnName: String, final val estimateName: String, sortedSizeMap: java.util.SortedMap[java.lang.Long, java.lang.Double]) extends Aggregator[Row, BloomFilter, Array[Byte]]
   with Serializable {
 
   var tokenizer: Option[Tokenizer] = None
@@ -113,14 +114,19 @@ class BloomFilterAggregator(final val columnName: String, final val estimateName
   implicit def customKryoEncoder[A](implicit ct: ClassTag[A]): Encoder[A] = Encoders.kryo[A](ct)
 
   private def selectFilterFromMap(estimate: Long): BloomFilter = {
-    var filter = BloomFilter.create(sizeMap.last._1, sizeMap.last._2)
+    val sortedScalaMap = sortedSizeMap.asScala
 
-    for (entry <- sizeMap) {
-      if (estimate <= entry._1) {
-        filter = BloomFilter.create(entry._1, entry._2)
+    // default to largest
+    var size = sortedScalaMap.last._1
+
+    for (entry <- sortedScalaMap) {
+      if (entry._1 >= estimate && entry._1 < size) {
+        size = entry._1
       }
     }
+    val fpp = sortedScalaMap.getOrElse(size,
+      throw new RuntimeException("sortedScalaMap did not contain value for key size: " + size))
 
-    filter
+    BloomFilter.create(size, fpp)
   }
 }
diff --git a/src/test/scala/BloomFilterAggregatorTest.scala b/src/test/scala/BloomFilterAggregatorTest.scala
@@ -55,7 +55,6 @@ import org.apache.spark.util.sketch.BloomFilter
 import java.io.ByteArrayInputStream
 import java.sql.Timestamp
 import java.time.{Instant, LocalDateTime, ZoneOffset}
-import scala.collection.mutable
 import scala.collection.mutable.ArrayBuffer
 
 class BloomFilterAggregatorTest {
@@ -91,7 +90,10 @@ class BloomFilterAggregatorTest {
     val rowMemoryStream = new MemoryStream[Row](1,sqlContext)(encoder)
 
     var rowDataset = rowMemoryStream.toDF
-    val sizeMap: mutable.TreeMap[Long, Double] = mutable.TreeMap(1000L -> 0.01, 10000L -> 0.01)
+    val javaMap = new java.util.TreeMap[java.lang.Long, java.lang.Double]() {
+      put(1000L, 0.01)
+      put(10000L, 0.01)
+    }
 
 
     // create Scala udf
@@ -103,7 +105,7 @@ class BloomFilterAggregatorTest {
     rowDataset = rowDataset.withColumn("tokens", tokenizerUDF.apply(functions.col("_raw")))
 
     // run bloomfilter on the column
-    val tokenAggregator = new BloomFilterAggregator("tokens", "estimate(tokens)", sizeMap)
+    val tokenAggregator = new BloomFilterAggregator("tokens", "estimate(tokens)", javaMap)
     val tokenAggregatorColumn = tokenAggregator.toColumn
 
     val aggregatedDataset = rowDataset

diff --git a/src/test/scala/BloomFilterBufferTest.scala b/src/test/scala/BloomFilterBufferTest.scala
@@ -46,7 +46,7 @@
 
 import com.teragrep.functions.dpf_03.BloomFilterAggregator
 import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
-import org.apache.spark.sql.types.{ArrayType, ByteType, StringType, StructField, StructType}
+import org.apache.spark.sql.types.{ArrayType, ByteType, StructField, StructType}
 import org.apache.spark.util.sketch.BloomFilter
 import org.junit.jupiter.api.Disabled
 
@@ -61,7 +61,10 @@ class BloomFilterBufferTest {
   def testNoDuplicateKeys(): Unit = {
 
     // TODO test other sizes / size categorization
-    val sizeMap: mutable.TreeMap[Long, Double] = mutable.TreeMap(1000L -> 0.01, 10000L -> 0.01)
+    val javaMap = new java.util.TreeMap[java.lang.Long, java.lang.Double]() {
+      put(1000L, 0.01)
+      put(10000L, 0.01)
+    }
 
     // single token, converted to WrappedArray
     val input: String = "one,one"
@@ -79,7 +82,7 @@ class BloomFilterBufferTest {
     val schema = StructType(Seq(StructField(columnName, ArrayType(ArrayType(ByteType)))))
     val row = new GenericRowWithSchema(columns, schema)
 
-    val bfAgg : BloomFilterAggregator = new BloomFilterAggregator(columnName, "estimate(tokens)", sizeMap)
+    val bfAgg : BloomFilterAggregator = new BloomFilterAggregator(columnName, "estimate(tokens)", javaMap)
 
     val bfAggBuf = bfAgg.zero()
     bfAgg.reduce(bfAggBuf, row)