repackage

hnswlib-spark/src/main/scala/com/github/jelmerk/server/registration/RegistrationClient.scala → hnswlib-spark/src/main/scala/com/github/jelmerk/registration/client/RegistrationClient.scala

@@ -1,10 +1,11 @@
-package com.github.jelmerk.server.registration
+package com.github.jelmerk.registration.client
 
 import java.net.{InetSocketAddress, SocketAddress}
 
 import scala.concurrent.Await
 import scala.concurrent.duration.Duration
 
+import com.github.jelmerk.server.registration.{RegisterRequest, RegisterResponse, RegistrationServiceGrpc}
 import io.grpc.netty.NettyChannelBuilder
 
 object RegistrationClient {

hnswlib-spark/src/main/scala/com/github/jelmerk/server/registration/DefaultRegistrationService.scala → hnswlib-spark/src/main/scala/com/github/jelmerk/registration/server/DefaultRegistrationService.scala

@@ -1,10 +1,11 @@
-package com.github.jelmerk.server.registration
+package com.github.jelmerk.registration.server
 
 import java.net.InetSocketAddress
 import java.util.concurrent.{ConcurrentHashMap, CountDownLatch}
 
 import scala.concurrent.Future
 
+import com.github.jelmerk.server.registration.{RegisterRequest, RegisterResponse}
 import com.github.jelmerk.server.registration.RegistrationServiceGrpc.RegistrationService
 
 class DefaultRegistrationService(val registrationLatch: CountDownLatch) extends RegistrationService {

hnswlib-spark/src/main/scala/com/github/jelmerk/server/registration/RegistrationServerFactory.scala → hnswlib-spark/src/main/scala/com/github/jelmerk/registration/server/RegistrationServerFactory.scala

@@ -1,4 +1,4 @@
-package com.github.jelmerk.server.registration
+package com.github.jelmerk.registration.server
 
 import java.net.{InetAddress, InetSocketAddress}
 import java.util.concurrent.{CountDownLatch, Executors}
@@ -7,6 +7,7 @@ import scala.concurrent.ExecutionContext
 import scala.jdk.CollectionConverters.mapAsScalaConcurrentMapConverter
 import scala.util.Try
 
+import com.github.jelmerk.server.registration.RegistrationServiceGrpc
 import io.grpc.netty.NettyServerBuilder
 
 // TODO this is all a bit messy

hnswlib-spark/src/main/scala/com/github/jelmerk/spark/knn/IndexClient.scala → hnswlib-spark/src/main/scala/com/github/jelmerk/serving/client/IndexClient.scala

@@ -1,4 +1,4 @@
-package com.github.jelmerk.spark.knn
+package com.github.jelmerk.serving.client
 
 import java.net.InetSocketAddress
 import java.util.concurrent.{Executors, LinkedBlockingQueue}
@@ -10,8 +10,8 @@ import scala.concurrent.duration.Duration
 import scala.language.implicitConversions
 import scala.util.Random
 
+import com.github.jelmerk.registration.server.PartitionAndReplica
 import com.github.jelmerk.server.index._
-import com.github.jelmerk.server.registration.PartitionAndReplica
 import io.grpc.netty.NettyChannelBuilder
 import io.grpc.stub.StreamObserver
 

hnswlib-spark/src/main/scala/com/github/jelmerk/spark/knn/IndexClientFactory.scala → hnswlib-spark/src/main/scala/com/github/jelmerk/serving/client/IndexClientFactory.scala

@@ -1,9 +1,9 @@
-package com.github.jelmerk.spark.knn
+package com.github.jelmerk.serving.client
 
 import java.net.InetSocketAddress
 
+import com.github.jelmerk.registration.server.PartitionAndReplica
 import com.github.jelmerk.server.index.{Result, SearchRequest}
-import com.github.jelmerk.server.registration.PartitionAndReplica
 
 class IndexClientFactory[TId, TVector, TDistance](
     vectorConverter: TVector => SearchRequest.Vector,

hnswlib-spark/src/main/scala/com/github/jelmerk/serving/client/Neighbor.scala

@@ -0,0 +1,15 @@
+package com.github.jelmerk.serving.client
+
+/** Neighbor of an item.
+  *
+  * @param neighbor
+  *   identifies the neighbor
+  * @param distance
+  *   distance to the item
+  *
+  * @tparam TId
+  *   type of the index item identifier
+  * @tparam TDistance
+  *   type of distance
+  */
+case class Neighbor[TId, TDistance](neighbor: TId, distance: TDistance) // TODO move this

hnswlib-spark/src/main/scala/com/github/jelmerk/server/index/DefaultIndexService.scala → hnswlib-spark/src/main/scala/com/github/jelmerk/serving/server/DefaultIndexService.scala

@@ -1,9 +1,10 @@
-package com.github.jelmerk.server.index
+package com.github.jelmerk.serving.server
 
 import scala.concurrent.{ExecutionContext, Future}
 import scala.language.implicitConversions
 
 import com.github.jelmerk.knn.scalalike.{Index, Item}
+import com.github.jelmerk.server.index._
 import com.github.jelmerk.server.index.IndexServiceGrpc.IndexService
 import io.grpc.stub.StreamObserver
 import org.apache.commons.io.output.CountingOutputStream

hnswlib-spark/src/main/scala/com/github/jelmerk/server/index/IndexServerFactory.scala → hnswlib-spark/src/main/scala/com/github/jelmerk/serving/server/IndexServerFactory.scala

@@ -1,4 +1,4 @@
-package com.github.jelmerk.server.index
+package com.github.jelmerk.serving.server
 
 import java.net.{InetAddress, InetSocketAddress}
 import java.util.concurrent.{LinkedBlockingQueue, ThreadPoolExecutor, TimeUnit}
@@ -7,6 +7,7 @@ import scala.concurrent.ExecutionContext
 import scala.util.Try
 
 import com.github.jelmerk.knn.scalalike.{Index, Item}
+import com.github.jelmerk.server.index.{IndexServiceGrpc, Result, SearchRequest}
 import io.grpc.netty.NettyServerBuilder
 import org.apache.hadoop.conf.Configuration
 

hnswlib-spark/src/main/scala/com/github/jelmerk/spark/knn/KnnAlgorithm.scala

@@ -10,8 +10,10 @@ import scala.util.control.NonFatal
 
 import com.github.jelmerk.knn.ObjectSerializer
 import com.github.jelmerk.knn.scalalike._
-import com.github.jelmerk.server.index.IndexServerFactory
-import com.github.jelmerk.server.registration.{PartitionAndReplica, RegistrationClient, RegistrationServerFactory}
+import com.github.jelmerk.registration.client.RegistrationClient
+import com.github.jelmerk.registration.server.{PartitionAndReplica, RegistrationServerFactory}
+import com.github.jelmerk.serving.client.IndexClientFactory
+import com.github.jelmerk.serving.server.IndexServerFactory
 import com.github.jelmerk.spark.util.SerializableConfiguration
 import org.apache.hadoop.fs.Path
 import org.apache.spark.{Partitioner, SparkContext, TaskContext}
@@ -43,20 +45,6 @@ private[knn] case class ModelMetaData(
     paramMap: Map[String, Any]
 )
 
-/** Neighbor of an item.
-  *
-  * @param neighbor
-  *   identifies the neighbor
-  * @param distance
-  *   distance to the item
-  *
-  * @tparam TId
-  *   type of the index item identifier
-  * @tparam TDistance
-  *   type of distance
-  */
-private[knn] case class Neighbor[TId, TDistance](neighbor: TId, distance: TDistance)
-
 private[knn] trait IndexType {
 
   /** Type of index. */
@@ -501,7 +489,7 @@ private[knn] abstract class KnnModelReader[TModel <: KnnModelBase[TModel]: Class
 
     val indexRdd = sc
       .makeRDD(partitionPaths)
-      .partitionBy(new PartitionIdPassthrough(metadata.numPartitions))
+      .partitionBy(new PartitionIdPartitioner(metadata.numPartitions))
       .withResources(profile)
       .mapPartitions { it =>
         val (partitionId, indexPath) = it.next()
@@ -791,7 +779,7 @@ private[knn] abstract class KnnAlgorithm[TModel <: KnnModelBase[TModel]](overrid
           )
           .as[(Int, TItem)]
           .rdd
-          .partitionBy(new PartitionIdPassthrough(getNumPartitions))
+          .partitionBy(new PartitionIdPartitioner(getNumPartitions))
           .values
       else
         dataset
@@ -886,17 +874,22 @@ private[knn] abstract class KnnAlgorithm[TModel <: KnnModelBase[TModel]](overrid
 
 }
 
-/** Partitioner that uses precomputed partitions
+/** Partitioner that uses precomputed partitions. Each partition id is its own partition
   *
   * @param numPartitions
   *   number of partitions
   */
-private[knn] class PartitionIdPassthrough(override val numPartitions: Int) extends Partitioner {
+private[knn] class PartitionIdPartitioner(override val numPartitions: Int) extends Partitioner {
   override def getPartition(key: Any): Int = key.asInstanceOf[Int]
 }
 
-// TODO rename this
-private[knn] class PartitionReplicaPartitioner(partitions: Int, replicas: Int) extends Partitioner {
+/**
+ * Partitioner that uses precomputed partitions. Each unique partition and replica combination is own partition
+ *
+ * @param partitions the total number of partitions
+ * @param replicas the total number of replicas
+ */
+private[knn] class PartitionReplicaIdPartitioner(partitions: Int, replicas: Int) extends Partitioner {
 
   override def numPartitions: Int = partitions + (replicas * partitions)
 
@@ -928,7 +921,9 @@ private[knn] trait ModelLogging extends Logging {
   protected def logInfo(partition: Int, message: String): Unit = println(f"partition $partition%04d: $message")
 
 //  protected def logInfo(partition: Int, replica: Int, message: String): Unit = logInfo(f"partition $partition%04d replica $partition%04d: $message")
-  protected def logInfo(partition: Int, replica: Int, message: String): Unit = println(f"partition $partition%04d replica $partition%04d: $message")
+  protected def logInfo(partition: Int, replica: Int, message: String): Unit = println(
+    f"partition $partition%04d replica $partition%04d: $message"
+  )
 }
 
 private[knn] trait IndexServing extends ModelLogging with IndexType {
@@ -956,7 +951,7 @@ private[knn] trait IndexServing extends ModelLogging with IndexType {
     }
 
     val replicaRdd =
-      if (numReplicas > 0) keyedIndexRdd.partitionBy(new PartitionReplicaPartitioner(numPartitions, numReplicas))
+      if (numReplicas > 0) keyedIndexRdd.partitionBy(new PartitionReplicaIdPartitioner(numPartitions, numReplicas))
       else keyedIndexRdd
 
     val server = RegistrationServerFactory.create(numPartitions, numReplicas)
@@ -974,7 +969,11 @@ private[knn] trait IndexServing extends ModelLogging with IndexType {
 
           val serverAddress = server.address
 
-          logInfo(partitionNum, replicaNum, s"started index server on host ${serverAddress.getHostName}:${serverAddress.getPort}")
+          logInfo(
+            partitionNum,
+            replicaNum,
+            s"started index server on host ${serverAddress.getHostName}:${serverAddress.getPort}"
+          )
 
           logInfo(
             partitionNum,

hnswlib-spark/src/main/scala/com/github/jelmerk/spark/knn/QueryIterator.scala

@@ -2,7 +2,8 @@ package com.github.jelmerk.spark.knn
 
 import java.net.InetSocketAddress
 
-import com.github.jelmerk.server.registration.PartitionAndReplica
+import com.github.jelmerk.registration.server.PartitionAndReplica
+import com.github.jelmerk.serving.client.{IndexClientFactory, Neighbor}
 
 class QueryIterator[TId, TVector, TDistance, TQueryId](
     indices: Map[PartitionAndReplica, InetSocketAddress],

hnswlib-spark/src/main/scala/com/github/jelmerk/spark/knn/bruteforce/BruteForceSimilarity.scala

@@ -9,7 +9,8 @@ import scala.reflect.runtime.universe._
 import com.github.jelmerk.knn.ObjectSerializer
 import com.github.jelmerk.knn.scalalike.{DistanceFunction, Item}
 import com.github.jelmerk.knn.scalalike.bruteforce.BruteForceIndex
-import com.github.jelmerk.server.registration.PartitionAndReplica
+import com.github.jelmerk.registration.server.PartitionAndReplica
+import com.github.jelmerk.serving.client.IndexClientFactory
 import com.github.jelmerk.spark.knn._
 import org.apache.spark.SparkContext
 import org.apache.spark.ml.param._

hnswlib-spark/src/main/scala/com/github/jelmerk/spark/knn/hnsw/HnswSimilarity.scala

@@ -9,7 +9,8 @@ import scala.reflect.runtime.universe._
 import com.github.jelmerk.knn
 import com.github.jelmerk.knn.scalalike.{DistanceFunction, Item}
 import com.github.jelmerk.knn.scalalike.hnsw._
-import com.github.jelmerk.server.registration.PartitionAndReplica
+import com.github.jelmerk.registration.server.PartitionAndReplica
+import com.github.jelmerk.serving.client.IndexClientFactory
 import com.github.jelmerk.spark.knn._
 import org.apache.spark.SparkContext
 import org.apache.spark.ml.param._

hnswlib-spark/src/main/scala/com/github/jelmerk/spark/knn/knn.scala

@@ -37,11 +37,12 @@ import com.github.jelmerk.server.index.{
   DenseVector,
   DoubleArrayVector,
   FloatArrayVector,
-  IndexServerFactory,
   Result,
   SearchRequest,
   SparseVector
 }
+import com.github.jelmerk.serving.client.IndexClientFactory
+import com.github.jelmerk.serving.server.IndexServerFactory
 import com.github.jelmerk.spark.linalg.functions.VectorDistanceFunctions
 import org.apache.spark.ml.linalg.{DenseVector => SparkDenseVector, SparseVector => SparkSparseVector, Vector, Vectors}
 

hnswlib-spark/src/test/scala/ServerTest.scala

@@ -4,7 +4,8 @@ import scala.concurrent.ExecutionContext
 
 import com.github.jelmerk.knn.scalalike.{floatCosineDistance, Item}
 import com.github.jelmerk.knn.scalalike.hnsw.HnswIndex
-import com.github.jelmerk.server.index.{DefaultIndexService, IndexServiceGrpc}
+import com.github.jelmerk.server.index.IndexServiceGrpc
+import com.github.jelmerk.serving.server.DefaultIndexService
 import io.grpc.netty.NettyServerBuilder
 import org.apache.hadoop.conf.Configuration