Ash/Selva - Added full and prefix scan support to all the stores

suuchi-core/src/main/scala/in/ashwanthkumar/suuchi/store/InMemoryStore.scala

@@ -1,19 +1,23 @@
 package in.ashwanthkumar.suuchi.store
 
 import java.nio.ByteBuffer
-import java.util.concurrent.ConcurrentHashMap
+import java.util.{Arrays => JArrays}
+import java.util.concurrent.ConcurrentSkipListMap
 
 import org.slf4j.LoggerFactory
 
+import scala.collection.JavaConversions._
+
 class InMemoryStore extends Store {
   private val log = LoggerFactory.getLogger(getClass)
-  private val store = new ConcurrentHashMap[ByteBuffer, Array[Byte]]()
+  private val store = new ConcurrentSkipListMap[ByteBuffer, Array[Byte]]()
 
   override def put(key: Array[Byte], value: Array[Byte]): Boolean = {
     log.trace(s"Put with key=${new String(key)}, value=${new String(value)}")
     store.put(ByteBuffer.wrap(key), value)
     true
   }
+
   override def get(key: Array[Byte]): Option[Array[Byte]] = {
     log.trace(s"Get with key=${new String(key)}")
     val value = Option(store.get(ByteBuffer.wrap(key)))
@@ -26,4 +30,19 @@ class InMemoryStore extends Store {
     store.remove(ByteBuffer.wrap(key))
     true
   }
+
+  override def scan(): Iterator[KV] = {
+    store.entrySet().map(kv => KV(kv.getKey.array(), kv.getValue)).iterator
+  }
+
+  private def hasPrefix(key: Array[Byte], prefix: Array[Byte]) = {
+    key.length >= prefix.length && JArrays.equals(key.take(prefix.length), prefix)
+  }
+
+  override def scan(prefix: Array[Byte]): Iterator[KV] = {
+    store.tailMap(ByteBuffer.wrap(prefix))
+      .takeWhile{case (k, v) => hasPrefix(k.array(), prefix)}
+      .map{case (k, v) => KV(k.array(), v)}
+      .iterator
+  }
 }

suuchi-core/src/main/scala/in/ashwanthkumar/suuchi/store/ShardedStore.scala

@@ -5,6 +5,8 @@ import java.util.concurrent.ConcurrentHashMap
 import in.ashwanthkumar.suuchi.partitioner.Hash
 import in.ashwanthkumar.suuchi.utils.Logging
 
+import scala.collection.JavaConversions._
+
 /**
  * SharedStore shards the keys equally into [[partitionsPerNode]] stores and proxies store operations
  * against them for a given key.
@@ -50,4 +52,8 @@ class ShardedStore(partitionsPerNode: Int, hashFn: Hash, createStore: (Int) => S
       }
     }
   }
+
+  override def scan(): Iterator[KV] = map.values().flatMap(_.scan()).iterator
+
+  override def scan(prefix: Array[Byte]): Iterator[KV] = map.values().flatMap(_.scan(prefix)).iterator
 }

suuchi-core/src/main/scala/in/ashwanthkumar/suuchi/store/Store.scala

@@ -1,5 +1,7 @@
 package in.ashwanthkumar.suuchi.store
 
+import java.util
+
 trait ReadStore {
   def get(key: Array[Byte]): Option[Array[Byte]]
 }
@@ -9,4 +11,17 @@ trait WriteStore {
   def remove(key: Array[Byte]): Boolean
 }
 
-trait Store extends ReadStore with WriteStore
+case class KV(key: Array[Byte], value: Array[Byte]) {
+  override def equals(obj: scala.Any): Boolean = obj match {
+    case o: KV => util.Arrays.equals(key, o.key) && util.Arrays.equals(value, o.value)
+    case _ => false
+  }
+
+  override def toString: String = new String(key) + " " + new String(value)
+}
+trait Scannable {
+  def scan(): Iterator[KV]
+  def scan(prefix: Array[Byte]): Iterator[KV]
+}
+
+trait Store extends ReadStore with WriteStore with Scannable

suuchi-core/src/main/scala/in/ashwanthkumar/suuchi/store/VersionedStore.scala

@@ -1,8 +1,10 @@
 package in.ashwanthkumar.suuchi.store
 
+import java.util
+
 import in.ashwanthkumar.suuchi.utils.DateUtils
-import scala.util.hashing.MurmurHash3
 
+import scala.util.hashing.MurmurHash3
 import scala.language.postfixOps
 
 
@@ -33,6 +35,7 @@ object VersionedStore {
   val VERSION_PREFIX = "V_".getBytes
   val DATA_PREFIX = "D_".getBytes
 
+  def isDataKey(key: Array[Byte]) = util.Arrays.equals(key.take(DATA_PREFIX.length), DATA_PREFIX)
   def vkey(key: Array[Byte]) = VERSION_PREFIX ++ key
   def dkey(key: Array[Byte], version: Array[Byte]): Array[Byte] = DATA_PREFIX ++ key ++ version
   def dkey(key: Array[Byte], version: Long): Array[Byte] = DATA_PREFIX ++ key ++ PrimitivesSerDeUtils.longToBytes(version)
@@ -107,4 +110,8 @@ class VersionedStore(store: Store, versionedBy: VersionedBy, numVersions: Int, c
 
   private def removeData(key: Array[Byte], version: Long) = store.remove(dkey(key, version))
   private def removeVersion(key: Array[Byte]) = store.remove(vkey(key))
+
+  override def scan(): Iterator[KV] = store.scan().filter(kv => VersionedStore.isDataKey(kv.key))
+
+  override def scan(prefix: Array[Byte]): Iterator[KV] = store.scan(prefix).filter(kv => VersionedStore.isDataKey(kv.key))
 }

suuchi-core/src/test/scala/in/ashwanthkumar/suuchi/store/InMemoryStoreTest.scala

@@ -1,12 +1,43 @@
 package in.ashwanthkumar.suuchi.store
 
 import org.scalatest.FlatSpec
-import org.scalatest.Matchers.{be, convertToAnyShouldWrapper}
+import org.scalatest.Matchers.{be, convertToAnyShouldWrapper, have, startWith}
 
 class InMemoryStoreTest extends FlatSpec {
   "InMemoryStore" should "support get and put on a KV" in {
     val store = new InMemoryStore
     store.put("1".getBytes, "2".getBytes) should be(true)
     store.get("1".getBytes).map(v => new String(v)) should be(Some("2"))
   }
+
+  it should "support full store scan" in {
+    val store = new InMemoryStore
+    store.put("1".getBytes, "one".getBytes)
+    store.put("2".getBytes, "two".getBytes)
+    store.put("3".getBytes, "three".getBytes)
+    store.put("4".getBytes, "four".getBytes)
+    store.put("5".getBytes, "five".getBytes)
+
+    val kVs = store.scan().toList
+    kVs should have size 5
+    kVs.sortBy(kv => new String(kv.key)) should be(List(kv("1", "one"), kv("2", "two"), kv("3", "three"), kv("4", "four"), kv("5", "five")))
+  }
+
+  it should "support prefix scan" in {
+    val store = new InMemoryStore
+    store.put("prefix1/1".getBytes, "one".getBytes)
+    store.put("prefix1/2".getBytes, "two".getBytes)
+    store.put("prefix1/3".getBytes, "three".getBytes)
+    store.put("prefix2/1".getBytes, "eleven".getBytes)
+    store.put("prefix2/2".getBytes, "twelve".getBytes)
+    store.put("prefix2/3".getBytes, "thirteen".getBytes)
+
+    val kVs = store.scan("prefix1".getBytes).toList
+
+    kVs.foreach{kv =>
+      new String(kv.key) should startWith("prefix1")
+    }
+  }
+
+  def kv(key: String, value: String) = KV(key.getBytes, value.getBytes)
 }

suuchi-rocksdb/src/main/scala/in/ashwanthkumar/suuchi/store/rocksdb/RocksDbStore.scala

@@ -1,6 +1,8 @@
 package in.ashwanthkumar.suuchi.store.rocksdb
 
-import in.ashwanthkumar.suuchi.store.Store
+import java.util.{Arrays => JArrays}
+
+import in.ashwanthkumar.suuchi.store.{KV, Store}
 import in.ashwanthkumar.suuchi.utils.Logging
 import org.rocksdb._
 
@@ -31,4 +33,26 @@ class RocksDbStore(config: RocksDbConfiguration) extends Store with Logging {
   override def remove(key: Array[Byte]): Boolean = {
     logOnError(() => db.remove(key)) isSuccess
   }
+
+  def scan(): Iterator[KV] = scan(Array.ofDim[Byte](0))
+
+  def scan(prefix: Array[Byte]): Iterator[KV] = {
+    val rocksIterator: RocksIterator = db.newIterator()
+    rocksIterator.seek(prefix)
+
+    val iterator = new Iterator[KV] {
+      override def hasNext: Boolean = rocksIterator.isValid && hasPrefix(rocksIterator.key(), prefix)
+      override def next(): KV = {
+        val kv = KV(rocksIterator.key(), rocksIterator.value())
+        rocksIterator.next()
+        kv
+      }
+
+      private def hasPrefix(key: Array[Byte], prefix: Array[Byte]) = {
+        key.length >= prefix.length && JArrays.equals(key.take(prefix.length), prefix)
+      }
+    }
+    iterator
+  }
+
 }

suuchi-rocksdb/src/test/scala/in/ashwanthkumar/suuchi/store/rocksdb/RocksDbStoreSpec.scala

@@ -2,29 +2,56 @@ package in.ashwanthkumar.suuchi.store.rocksdb
 
 import java.io.File
 
+import in.ashwanthkumar.suuchi.store.KV
 import org.apache.commons.io.FileUtils
-import org.scalatest.Matchers._
+import org.scalatest.Matchers.{be, convertToAnyShouldWrapper, have, startWith}
 import org.scalatest.{BeforeAndAfter, FlatSpec}
 
 class RocksDbStoreSpec extends FlatSpec with BeforeAndAfter {
-  val ROCKSDB_TEST_LOCATION = "/tmp/suuchi-rocks-test"
+
+  val ROCKSDB_TEST_LOCATION = "/tmp/suuchi-rocks-test/"
+
+  var db: RocksDbStore = _
+
   before {
     FileUtils.deleteDirectory(new File(ROCKSDB_TEST_LOCATION))
+    db = new RocksDbStore(RocksDbConfiguration.apply(ROCKSDB_TEST_LOCATION))
   }
 
   after {
     FileUtils.deleteDirectory(new File(ROCKSDB_TEST_LOCATION))
+    db.close()
   }
 
   "RocksDb" should "store & retrieve results properly" in {
-    val db = new RocksDbStore(RocksDbConfiguration.apply(ROCKSDB_TEST_LOCATION))
     (1 to 100).foreach { i =>
       db.put(Array(i toByte), Array(i*2 toByte))
     }
 
     (1 to 100).foreach { i =>
       db.get(Array(i toByte)).get.head should be(i*2 toByte)
     }
+  }
+
+  it should "support full db scan" in {
+    val inputKVs = (1 to 100).map(i => (Array(i toByte), Array(i*2 toByte)))
+
+    inputKVs.foreach{case (k, v) => db.put(k, v)}
+    val scannedResult = db.scan().toList
 
+    scannedResult should have size 100
+    scannedResult.sortBy(kv => new String(kv.key)) should be(inputKVs.map{case (k,v) => KV(k, v)}.toList)
+  }
+
+  it should "support prefix scan" in {
+    val kVs = (1 to 100).flatMap(i => List((s"prefix1/$i".getBytes, Array(i toByte)), (s"prefix2/$i".getBytes, Array(i * 2 toByte))))
+
+    kVs.foreach{case (k, v) => db.put(k, v)}
+    val scannedResult = db.scan("prefix1".getBytes).toList
+
+    scannedResult should have size 100
+    scannedResult.foreach{ kv =>
+      new String(kv.key) should startWith("prefix1")
+    }
   }
 }