Add Cassandra 5 SAI index and vector search support

.github/workflows/ci.yml

@@ -14,7 +14,7 @@ jobs:
       matrix:
         java: [ '8', '11', '17', '21' ]
         scala: [ '2.12.20', '2.13.16' ]
-        cassandra: [ '3.11', '4.0', '4.1', '5.0' ]
+        cassandra: [ '3.11', '4.0', '4.1', '5.0.7' ]
 
     name: Test JDK ${{ matrix.java }}, Scala ${{ matrix.scala }}, Cassandra ${{ matrix.cassandra }}
 

core/src/main/scala/spinoco/fs2/cassandra/builder/CollectionIndexTarget.scala

@@ -0,0 +1,12 @@
+package spinoco.fs2.cassandra.builder
+
+sealed trait CollectionIndexTarget {
+  def wrap(field: String): String
+}
+
+object CollectionIndexTarget {
+  case object Keys    extends CollectionIndexTarget { def wrap(f: String): String = s"KEYS($f)" }
+  case object Values  extends CollectionIndexTarget { def wrap(f: String): String = s"VALUES($f)" }
+  case object Entries extends CollectionIndexTarget { def wrap(f: String): String = s"ENTRIES($f)" }
+  case object Full    extends CollectionIndexTarget { def wrap(f: String): String = s"FULL($f)" }
+}

core/src/main/scala/spinoco/fs2/cassandra/builder/IndexEntry.scala

@@ -5,16 +5,18 @@ case class IndexEntry(
   , field: String
   , className: Option[String]
   , options: Map[String, String]
+  , collectionTarget: Option[CollectionIndexTarget] = None
 ) {
   def cqlStatement(ks: String, table: String): String = {
+    val fieldExpr = collectionTarget.fold(field)(_.wrap(field))
     className match {
-      case None => s"CREATE INDEX $name ON $ks.$table ($field)"
+      case None => s"CREATE INDEX $name ON $ks.$table ($fieldExpr)"
       case Some(clz) =>
         val withOptions =
           if (options.isEmpty) ""
           else options.map { case (k,v) => s"'$k': '$v'"}.mkString(" WITH OPTIONS = {",",","}")
 
-        s"CREATE CUSTOM INDEX $name ON $ks.$table ($field) USING '$clz'$withOptions"
+        s"CREATE CUSTOM INDEX $name ON $ks.$table ($fieldExpr) USING '$clz'$withOptions"
 
     }
   }
@@ -24,4 +26,5 @@ case class IndexEntry(
 object IndexEntry {
 
   val SASIIndexClz = "org.apache.cassandra.index.sasi.SASIIndex"
+  val SAIIndexClz  = "org.apache.cassandra.index.sai.StorageAttachedIndex"
 }

core/src/main/scala/spinoco/fs2/cassandra/builder/QueryBuilder.scala

@@ -8,6 +8,7 @@ import shapeless.ops.record.{Keys, Selector}
 import shapeless.{::, HList, HNil, Witness}
 import spinoco.fs2.cassandra._
 import spinoco.fs2.cassandra.baseutil.{AnnotatedException, replaceInCql}
+import spinoco.fs2.cassandra.ctype.CType
 import spinoco.fs2.cassandra.internal._
 import spinoco.fs2.cassandra.macros.CTypeRecord
 
@@ -21,6 +22,7 @@ case class QueryBuilder[R <: HList, PK <: HList, CK <: HList, IDX <: HList, Q <:
   , clusterColumns: Map[Comparison.Value, Seq[(String, String)]]
   , limitCount: Option[Int]
   , allowFilteringFlag: Boolean
+  , annOrderBy: Option[(String, String)] = None
 ) { self =>
 
   /** mark this query to contain all columns in the table as result **/
@@ -236,19 +238,148 @@ case class QueryBuilder[R <: HList, PK <: HList, CK <: HList, IDX <: HList, Q <:
   ):QueryBuilder[R, PK, CK, IDX,  Q, S] =
     copy( orderColumns = orderColumns :+ (internal.keyOf(name) -> ascending))
 
+  /** order results by approximate nearest neighbor on a vector column.
+    * The query vector is added to Q as a bound parameter.
+    * ANN queries require a LIMIT clause. */
+  def orderByAnn[K, V](
+    column: Witness.Aux[K]
+  )(implicit
+    ev0: Selector.Aux[IDX, K, V]
+    , P: Prepend[Q, FieldType[K, V] :: HNil]
+  ): QueryBuilder[R, PK, CK, IDX, P.Out, S] =
+    orderByAnn(column, column)
+
+  /** like `orderByAnn` but allows to specify alias for the bound parameter name */
+  def orderByAnn[K, K0, V](
+    column: Witness.Aux[K]
+    , as: Witness.Aux[K0]
+  )(implicit
+    ev0: Selector.Aux[IDX, K, V]
+    , P: Prepend[Q, FieldType[K0, V] :: HNil]
+  ): QueryBuilder[R, PK, CK, IDX, P.Out, S] = {
+    val k = internal.keyOf(column)
+    val k0 = internal.keyOf(as)
+    QueryBuilder(
+      table = table
+      , queryColumns = queryColumns
+      , whereConditions = whereConditions
+      , orderColumns = orderColumns
+      , clusterColumns = clusterColumns
+      , limitCount = limitCount
+      , allowFilteringFlag = allowFilteringFlag
+      , annOrderBy = Some((k, k0))
+    )
+  }
+
+  /** like `orderByAnn` but uses a parameter already present in Q.
+    * Use when the same bound parameter is shared between ANN and a similarity function. */
+  def orderByAnnShared[K, K0, V](
+    column: Witness.Aux[K]
+    , paramName: Witness.Aux[K0]
+  )(implicit
+    ev0: Selector.Aux[IDX, K, V]
+    , ev1: Selector.Aux[Q, K0, V]
+  ): QueryBuilder[R, PK, CK, IDX, Q, S] = {
+    val k = internal.keyOf(column)
+    val k0 = internal.keyOf(paramName)
+    copy(annOrderBy = Some((k, k0)))
+  }
+
+  /** returns only rows where indexed field is in the given list of values */
+  def byIndexIn[K, V](
+    column: Witness.Aux[K]
+  )(implicit
+    ev0: Selector.Aux[IDX, K, V]
+    , ev1: CType[List[V]]
+    , P: Prepend[Q, FieldType[K, List[V]] :: HNil]
+  ): QueryBuilder[R, PK, CK, IDX, P.Out, S] =
+    byIndexIn(column, column)
+
+  /** like `byIndexIn` but allows to specify alias for the bound parameter name */
+  def byIndexIn[K, K0, V](
+    column: Witness.Aux[K]
+    , as: Witness.Aux[K0]
+  )(implicit
+    ev0: Selector.Aux[IDX, K, V]
+    , ev1: CType[List[V]]
+    , P: Prepend[Q, FieldType[K0, List[V]] :: HNil]
+  ): QueryBuilder[R, PK, CK, IDX, P.Out, S] = {
+    val k = internal.keyOf(column)
+    val k0 = internal.keyOf(as)
+    QueryBuilder(
+      table = table
+      , queryColumns = queryColumns
+      , whereConditions = whereConditions :+ s"$k IN :$k0"
+      , orderColumns = orderColumns
+      , clusterColumns = clusterColumns
+      , limitCount = limitCount
+      , allowFilteringFlag = allowFilteringFlag
+      , annOrderBy = annOrderBy
+    )
+  }
+
+  /** select a two-argument function applied to a column and a bound parameter.
+    * The bound parameter is added to Q. */
+  def function2At[K, V, K0, V0, K1](
+    fn: CQLFunction2[V, V, V0]
+    , column: Witness.Aux[K]
+    , param: Witness.Aux[K0]
+    , as: Witness.Aux[K1]
+  )(implicit
+    ev0: Selector.Aux[R, K, V]
+    , P: Prepend[Q, FieldType[K0, V] :: HNil]
+  ): QueryBuilder[R, PK, CK, IDX, P.Out, FieldType[K1, V0] :: S] = {
+    QueryBuilder(
+      table = table
+      , queryColumns = queryColumns :+ (fn(internal.keyOf(column), internal.keyOf(param)) -> internal.keyOf(as))
+      , whereConditions = whereConditions
+      , orderColumns = orderColumns
+      , clusterColumns = clusterColumns
+      , limitCount = limitCount
+      , allowFilteringFlag = allowFilteringFlag
+      , annOrderBy = annOrderBy
+    )
+  }
+
+  /** like `function2At` but uses a parameter already present in Q.
+    * Use when the same bound parameter is shared between ANN and a similarity function. */
+  def function2AtShared[K, V, K0, V0, K1](
+    fn: CQLFunction2[V, V, V0]
+    , column: Witness.Aux[K]
+    , param: Witness.Aux[K0]
+    , as: Witness.Aux[K1]
+  )(implicit
+    ev0: Selector.Aux[R, K, V]
+    , ev1: Selector.Aux[Q, K0, V]
+  ): QueryBuilder[R, PK, CK, IDX, Q, FieldType[K1, V0] :: S] = {
+    QueryBuilder(
+      table = table
+      , queryColumns = queryColumns :+ (fn(internal.keyOf(column), internal.keyOf(param)) -> internal.keyOf(as))
+      , whereConditions = whereConditions
+      , orderColumns = orderColumns
+      , clusterColumns = clusterColumns
+      , limitCount = limitCount
+      , allowFilteringFlag = allowFilteringFlag
+      , annOrderBy = annOrderBy
+    )
+  }
+
   /** creates query, that may be used to perform CQL commands on connection **/
   def build(
    implicit
    CTQ: CTypeRecord[Q]
    , CTS: CTypeRecord[S]
   ): Query[Q, S] = {
-    val orderStmt = {
-      val ocs =
-        orderColumns.map {
-          case (k, asc) => s"$k ${if(asc) "ASC" else "DESC"}"
-        }.mkString(",")
-
-      if (ocs.nonEmpty) s"ORDER BY $ocs" else ""
+    val orderStmt = annOrderBy match {
+      case Some((col, param)) =>
+        s"ORDER BY $col ANN OF :$param"
+      case None =>
+        val ocs =
+          orderColumns.map {
+            case (k, asc) => s"$k ${if(asc) "ASC" else "DESC"}"
+          }.mkString(",")
+
+        if (ocs.nonEmpty) s"ORDER BY $ocs" else ""
     }
 
 

core/src/main/scala/spinoco/fs2/cassandra/builder/SimilarityFunction.scala

@@ -0,0 +1,9 @@
+package spinoco.fs2.cassandra.builder
+
+sealed abstract class SimilarityFunction(val name: String)
+
+object SimilarityFunction {
+  case object COSINE      extends SimilarityFunction("COSINE")
+  case object DOT_PRODUCT extends SimilarityFunction("DOT_PRODUCT")
+  case object EUCLIDEAN   extends SimilarityFunction("EUCLIDEAN")
+}

core/src/main/scala/spinoco/fs2/cassandra/builder/TableBuilder.scala

@@ -67,6 +67,26 @@ case class TableBuilder[R <: HList, PK <: HList, CK <: HList, IDX <: HList](
   ): TableBuilder[R, PK, CK, FieldType[K, V] :: IDX] =
     indexBy(column, name, Some(IndexEntry.SASIIndexClz), Map("mode" -> "SPARSE") ++ options)
 
+  /** create SAI (Storage Attached Index) on specified column **/
+  def indexBySAI[K, V](column: Witness.Aux[K], name: String, options: Map[String, String] = Map.empty)(
+    implicit S: Selector.Aux[R, K, V]
+  ): TableBuilder[R, PK, CK, FieldType[K, V] :: IDX] =
+    indexBy(column, name, Some(IndexEntry.SAIIndexClz), options)
+
+  /** create SAI index with specific similarity function for vector columns **/
+  def indexBySAIVector[K, V](column: Witness.Aux[K], name: String, similarity: SimilarityFunction = SimilarityFunction.COSINE, options: Map[String, String] = Map.empty)(
+    implicit S: Selector.Aux[R, K, V]
+  ): TableBuilder[R, PK, CK, FieldType[K, V] :: IDX] =
+    indexBy(column, name, Some(IndexEntry.SAIIndexClz), options + ("similarity_function" -> similarity.name))
+
+  /** create SAI index on collection column with specified target (KEYS/VALUES/ENTRIES/FULL) **/
+  def indexBySAICollection[K, V](column: Witness.Aux[K], name: String, target: CollectionIndexTarget, options: Map[String, String] = Map.empty)(
+    implicit S: Selector.Aux[R, K, V]
+  ): TableBuilder[R, PK, CK, FieldType[K, V] :: IDX] = {
+    val entry = IndexEntry(name, internal.keyOf(column), Some(IndexEntry.SAIIndexClz), options, Some(target))
+    TableBuilder(ks, entry +: indexes, partitionKeys, clusterKeys)
+  }
+
   def build(name: String, options: Map[String, String] = Map.empty)(
     implicit T: TableInstance[R, PK, CK, IDX]
   ): Table[R, PK, CK, IDX] = T.table(ks,name,options, self.indexes, self.partitionKeys, self.clusterKeys)

core/src/main/scala/spinoco/fs2/cassandra/comparison.scala

@@ -6,4 +6,6 @@ object Comparison extends Enumeration {
   val GTEQ = Value(">=")
   val LT = Value("<")
   val LTEQ = Value("<=")
+  val CONTAINS = Value("CONTAINS")
+  val CONTAINS_KEY = Value("CONTAINS KEY")
 }

core/src/main/scala/spinoco/fs2/cassandra/functions.scala

@@ -16,6 +16,15 @@ object functions {
 
   def writeTimeOfMicro[I : CType] : CQLFunction[I, Long] = CQLFunction(name => s"WRITETIME($name)")
   def ttlOf[I: CType] : CQLFunction[I, Option[FiniteDuration @@ TTL]] = CQLFunction(name => s"TTL($name)")
+
+  def similarityCosine[A: CType, T]: CQLFunction2[Vector[A] @@ T, Vector[A] @@ T, Float] =
+    CQLFunction2((col, param) => s"similarity_cosine($col, :$param)")
+
+  def similarityEuclidean[A: CType, T]: CQLFunction2[Vector[A] @@ T, Vector[A] @@ T, Float] =
+    CQLFunction2((col, param) => s"similarity_euclidean($col, :$param)")
+
+  def similarityDotProduct[A: CType, T]: CQLFunction2[Vector[A] @@ T, Vector[A] @@ T, Float] =
+    CQLFunction2((col, param) => s"similarity_dot_product($col, :$param)")
 }
 
 /** cql function taking column as parameter **/
@@ -28,6 +37,11 @@ trait CQLFunction0[O] {
   def apply(): String
 }
 
+/** CQL function taking a column and a bound parameter **/
+trait CQLFunction2[I1, I2, O] {
+  def apply(column: String, param: String): String
+}
+
 object CQLFunction0 {
   def apply[O](s: String): CQLFunction0[O] = new CQLFunction0[O] {
     def apply(): String = s
@@ -40,3 +54,10 @@ object CQLFunction {
     new CQLFunction[I, O] { def apply(s: String): String = f(s) }
 
 }
+
+object CQLFunction2 {
+
+  def apply[I1, I2, O](f: (String, String) => String): CQLFunction2[I1, I2, O] =
+    new CQLFunction2[I1, I2, O] { def apply(c: String, p: String): String = f(c, p) }
+
+}

core/src/main/scala/spinoco/fs2/cassandra/system/system.scala

@@ -2,6 +2,7 @@ package spinoco.fs2.cassandra
 
 import com.datastax.oss.driver.api.core.`type`.DataType
 import com.datastax.oss.driver.api.core.metadata.schema._
+import spinoco.fs2.cassandra.builder.IndexEntry
 
 import scala.jdk.CollectionConverters._
 
@@ -80,9 +81,72 @@ package object system {
           lazy val tableTemplate = s"ALTER TABLE $fullTableName"
           val cqlRemoved = removed.map { case (k, _) => s"$tableTemplate DROP $k" }
           val cqlAdded = added.map {case (k, tpe) => s"$tableTemplate ADD $k ${tpe.asCql(true, false)}"}
-          val res = cqlRemoved ++ cqlAdded
+
+          val indexStatements = migrateIndexes(desiredTable, current)
+
+          val res = cqlRemoved ++ cqlAdded ++ indexStatements
           res.toSeq
         }
     }
   }
+
+  /** compares desired indexes against current indexes, returns DROP/CREATE statements **/
+  def migrateIndexes(desiredTable: Table[_, _, _, _], current: TableMetadata): Seq[String] = {
+    val currentIndexes = current.getIndexes.asScala
+    val desiredIndexes = desiredTable.indexes
+
+    val desiredByName = desiredIndexes.map(idx => idx.name.toLowerCase -> idx).toMap
+
+    // indexes to drop: exist in current but not in desired, or exist but changed
+    val toDrop = currentIndexes.flatMap { case (cqlId, meta) =>
+      val name = cqlId.asInternal.toLowerCase
+      desiredByName.get(name) match {
+        case None =>
+          // index exists in C* but not desired - drop it
+          Some(s"DROP INDEX ${desiredTable.keySpaceName}.$name")
+        case Some(desired) =>
+          // index exists in both - check if it changed
+          if (!sameIndex(desired, meta, desiredTable.keySpaceName, desiredTable.name)) {
+            Some(s"DROP INDEX ${desiredTable.keySpaceName}.$name")
+          } else None
+      }
+    }
+
+    // indexes to create: not in current, or were dropped because they changed
+    val droppedNames = toDrop.map(_.split('.').last.trim.toLowerCase).toSet
+    val currentNames = currentIndexes.keys.map(_.asInternal.toLowerCase).toSet
+
+    val toCreate = desiredIndexes.flatMap { desired =>
+      val name = desired.name.toLowerCase
+      if (!currentNames.contains(name) || droppedNames.contains(name)) {
+        Some(desired.cqlStatement(desiredTable.keySpaceName, desiredTable.name))
+      } else None
+    }
+
+    (toDrop ++ toCreate).toSeq
+  }
+
+  /** checks whether a desired IndexEntry matches the current IndexMetadata **/
+  def sameIndex(desired: IndexEntry, current: IndexMetadata, ks: String, table: String): Boolean = {
+    val currentOptions = current.getOptions.asScala
+
+    // compare class name
+    val classMatches = desired.className match {
+      case None => current.getKind != IndexKind.CUSTOM
+      case Some(clz) => currentOptions.get("class_name").contains(clz)
+    }
+
+    // compare target column
+    val targetMatches = currentOptions.get("target").exists { target =>
+      val desiredField = desired.collectionTarget.fold(desired.field)(_.wrap(desired.field))
+      target.equalsIgnoreCase(desiredField)
+    }
+
+    // compare options (exclude internal keys like class_name and target)
+    val internalKeys = Set("class_name", "target")
+    val currentUserOptions = currentOptions.filterNot { case (k, _) => internalKeys.contains(k) }
+    val optionsMatch = desired.options == currentUserOptions
+
+    classMatches && targetMatches && optionsMatch
+  }
 }

scripts/start-cassandra.sh

@@ -8,7 +8,7 @@ VERSION=${1:-"3.11"}
 PORT=${2:-12000}
 
 # Supported versions
-SUPPORTED_VERSIONS=("3.11" "4.0" "4.1" "5.0")
+SUPPORTED_VERSIONS=("3.11" "4.0" "4.1" "5.0" "5.0.7")
 
 # Check if version is supported
 if [[ ! " ${SUPPORTED_VERSIONS[@]} " =~ " ${VERSION} " ]]; then