diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4241727..0233d5a 100644 --- a/.github/workflows/ci.yml +++ b/.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 }} diff --git a/core/src/main/scala/spinoco/fs2/cassandra/builder/CollectionIndexTarget.scala b/core/src/main/scala/spinoco/fs2/cassandra/builder/CollectionIndexTarget.scala new file mode 100644 index 0000000..f23f396 --- /dev/null +++ b/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)" } +} diff --git a/core/src/main/scala/spinoco/fs2/cassandra/builder/IndexEntry.scala b/core/src/main/scala/spinoco/fs2/cassandra/builder/IndexEntry.scala index b07f266..42703f2 100644 --- a/core/src/main/scala/spinoco/fs2/cassandra/builder/IndexEntry.scala +++ b/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" } diff --git a/core/src/main/scala/spinoco/fs2/cassandra/builder/QueryBuilder.scala b/core/src/main/scala/spinoco/fs2/cassandra/builder/QueryBuilder.scala index e0ac4df..2c0c1b5 100644 --- a/core/src/main/scala/spinoco/fs2/cassandra/builder/QueryBuilder.scala +++ b/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 "" } diff --git a/core/src/main/scala/spinoco/fs2/cassandra/builder/SimilarityFunction.scala b/core/src/main/scala/spinoco/fs2/cassandra/builder/SimilarityFunction.scala new file mode 100644 index 0000000..6262c3f --- /dev/null +++ b/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") +} diff --git a/core/src/main/scala/spinoco/fs2/cassandra/builder/TableBuilder.scala b/core/src/main/scala/spinoco/fs2/cassandra/builder/TableBuilder.scala index d2ce2e6..63669e8 100644 --- a/core/src/main/scala/spinoco/fs2/cassandra/builder/TableBuilder.scala +++ b/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) diff --git a/core/src/main/scala/spinoco/fs2/cassandra/comparison.scala b/core/src/main/scala/spinoco/fs2/cassandra/comparison.scala index 5e9390f..92e3c07 100644 --- a/core/src/main/scala/spinoco/fs2/cassandra/comparison.scala +++ b/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") } \ No newline at end of file diff --git a/core/src/main/scala/spinoco/fs2/cassandra/functions.scala b/core/src/main/scala/spinoco/fs2/cassandra/functions.scala index 11d5e38..bcceb92 100644 --- a/core/src/main/scala/spinoco/fs2/cassandra/functions.scala +++ b/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) } + +} diff --git a/core/src/main/scala/spinoco/fs2/cassandra/system/system.scala b/core/src/main/scala/spinoco/fs2/cassandra/system/system.scala index 9e017b6..5f9c01e 100644 --- a/core/src/main/scala/spinoco/fs2/cassandra/system/system.scala +++ b/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 + } } diff --git a/scripts/start-cassandra.sh b/scripts/start-cassandra.sh index ee4a6f8..1f0c59c 100755 --- a/scripts/start-cassandra.sh +++ b/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 diff --git a/test-support/src/main/scala/spinoco/fs2/cassandra/support/DockerCassandra.scala b/test-support/src/main/scala/spinoco/fs2/cassandra/support/DockerCassandra.scala index f141940..6f1a9ec 100644 --- a/test-support/src/main/scala/spinoco/fs2/cassandra/support/DockerCassandra.scala +++ b/test-support/src/main/scala/spinoco/fs2/cassandra/support/DockerCassandra.scala @@ -62,7 +62,7 @@ trait DockerCassandra override protected def beforeAll(): Unit = { super.beforeAll() - // Assume Cassandra is already running (started externally) + // Assume Cassandra is already running (started externally via scripts/start-cassandra.sh) println(s"Connecting to Cassandra $cassandraVersion at 127.0.0.1:$cqlPort") val session = clusterConfig.build() val cs = CassandraSession.impl.mkSession[IO](session, session.getContext.getProtocolVersion).unsafeRunSync() @@ -72,7 +72,7 @@ trait DockerCassandra override protected def afterAll(): Unit = { sessionInstance.foreach(_._1.close()) - // NOTE: Container cleanup is handled externally + // NOTE: Container cleanup is handled externally via scripts/stop-cassandra.sh super.afterAll() } diff --git a/test/src/test/scala/spinoco/fs2/cassandra/MigrationsSpec.scala b/test/src/test/scala/spinoco/fs2/cassandra/MigrationsSpec.scala index 9520627..5cabf47 100644 --- a/test/src/test/scala/spinoco/fs2/cassandra/MigrationsSpec.scala +++ b/test/src/test/scala/spinoco/fs2/cassandra/MigrationsSpec.scala @@ -114,5 +114,62 @@ class MigrationsSpec extends SchemaSupport { ) } + "will add SAI index if missing" in withSessionFor(_.startsWith("5")) { cs => + val tableNoIdx = ks.table[FooTable1].partition(Symbol("intColumn")).build("foo_sai") + val tableWithIdx = ks.table[FooTable1].partition(Symbol("intColumn")) + .indexBySAI(Symbol("strColumn"), "str_sai_idx") + .build("foo_sai") + + cs.create(ks).unsafeRunSync() + cs.create(tableNoIdx).unsafeRunSync() + + val migrate = cs.migrateDDL(tableWithIdx).unsafeRunSync() + migrate shouldBe Seq( + "CREATE CUSTOM INDEX str_sai_idx ON crud_ks.foo_sai (strColumn) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex'" + ) + + // apply migration + migrate.foreach(cs.executeCql(_).unsafeRunSync()) + + // no further migration needed + cs.migrateDDL(tableWithIdx).unsafeRunSync() shouldBe Nil + } + + "will drop removed SAI index" in withSessionFor(_.startsWith("5")) { cs => + val tableWithIdx = ks.table[FooTable1].partition(Symbol("intColumn")) + .indexBySAI(Symbol("strColumn"), "str_sai_idx2") + .build("foo_sai2") + val tableNoIdx = ks.table[FooTable1].partition(Symbol("intColumn")).build("foo_sai2") + + cs.create(ks).unsafeRunSync() + cs.create(tableWithIdx).unsafeRunSync() + + val migrate = cs.migrateDDL(tableNoIdx).unsafeRunSync() + migrate shouldBe Seq("DROP INDEX crud_ks.str_sai_idx2") + + migrate.foreach(cs.executeCql(_).unsafeRunSync()) + cs.migrateDDL(tableNoIdx).unsafeRunSync() shouldBe Nil + } + + "will recreate SAI index when options change" in withSessionFor(_.startsWith("5")) { cs => + val tableV1 = ks.table[FooTable1].partition(Symbol("intColumn")) + .indexBySAI(Symbol("strColumn"), "str_sai_idx3", Map("case_sensitive" -> "true")) + .build("foo_sai3") + val tableV2 = ks.table[FooTable1].partition(Symbol("intColumn")) + .indexBySAI(Symbol("strColumn"), "str_sai_idx3", Map("case_sensitive" -> "false")) + .build("foo_sai3") + + cs.create(ks).unsafeRunSync() + cs.create(tableV1).unsafeRunSync() + + val migrate = cs.migrateDDL(tableV2).unsafeRunSync() + migrate.size shouldBe 2 + migrate.head shouldBe "DROP INDEX crud_ks.str_sai_idx3" + migrate(1) should include("case_sensitive") + + migrate.foreach(cs.executeCql(_).unsafeRunSync()) + cs.migrateDDL(tableV2).unsafeRunSync() shouldBe Nil + } + } } diff --git a/test/src/test/scala/spinoco/fs2/cassandra/SAISpec.scala b/test/src/test/scala/spinoco/fs2/cassandra/SAISpec.scala new file mode 100644 index 0000000..5cd73ab --- /dev/null +++ b/test/src/test/scala/spinoco/fs2/cassandra/SAISpec.scala @@ -0,0 +1,299 @@ +package spinoco.fs2.cassandra + +import fs2.Stream._ +import shapeless.tag +import shapeless.tag.@@ +import spinoco.fs2.cassandra.builder.SimilarityFunction +import spinoco.fs2.cassandra.sample.VectorSizes._ +import spinoco.fs2.cassandra.sample.{SimpleTableRow, VectorTableRow} + + +class SAISpec extends SchemaSupport { + + def withCassandra5(f: CassandraSession[cats.effect.IO] => Any): Unit = + withSessionFor(_.startsWith("5"))(f) + + "SAI index features" - { + + "create table with SAI index" in withCassandra5 { cs => + val table = + ks.table[SimpleTableRow] + .partition(Symbol("intColumn")) + .cluster(Symbol("longColumn")) + .indexBySAI(Symbol("asciiColumn"), "ascii_sai_idx") + .build("sai_simple_table") + + (for { + _ <- cs.create(ks) + _ <- cs.create(table) + } yield ()).unsafeRunSync() + + val query = system.schema.queryAllTables.map(t => t.keyspace_name -> t.table_name) + val result = cs.queryAll(query).compile.toVector.unsafeRunSync() + result should contain(ks.name -> "sai_simple_table") + } + + "query by SAI index with EQ" in withCassandra5 { cs => + val table = + ks.table[SimpleTableRow] + .partition(Symbol("intColumn")) + .cluster(Symbol("longColumn")) + .indexBySAI(Symbol("stringColumn"), "string_sai_idx") + .build("sai_eq_table") + + (for { + _ <- cs.create(ks) + _ <- cs.create(table) + } yield ()).unsafeRunSync() + + val insert = table.insert.all.build.from[SimpleTableRow] + val entries = for (i <- 0 to 5; l <- 0L to 2L) yield + SimpleTableRow.simpleInstance.copy( + intColumn = i, longColumn = l, + stringColumn = if (i % 2 == 0) "even" else "odd" + ) + emits(entries).flatMap(e => eval(cs.execute(insert)(e)).drain).compile.drain.unsafeRunSync() + + val query = table.query.all + .byIndex(Symbol("stringColumn"), Comparison.EQ) + .allowFiltering + .build + .fromA[String] + .as[SimpleTableRow] + + val result = cs.query(query)("even").compile.toVector.unsafeRunSync() + result.foreach(_.stringColumn shouldBe "even") + result should not be empty + } + + "query by SAI index with range" in withCassandra5 { cs => + val table = + ks.table[SimpleTableRow] + .partition(Symbol("intColumn")) + .cluster(Symbol("longColumn")) + .indexBySAI(Symbol("floatColumn"), "float_sai_idx") + .build("sai_range_table") + + (for { + _ <- cs.create(ks) + _ <- cs.create(table) + } yield ()).unsafeRunSync() + + val insert = table.insert.all.build.from[SimpleTableRow] + val entries = for (i <- 0 to 3) yield + SimpleTableRow.simpleInstance.copy(intColumn = i, longColumn = 1L, floatColumn = i.toFloat * 10.0f) + emits(entries).flatMap(e => eval(cs.execute(insert)(e)).drain).compile.drain.unsafeRunSync() + + val query = table.query.all + .byIndex(Symbol("floatColumn"), Comparison.GTEQ) + .allowFiltering + .build + .fromA[Float] + .as[SimpleTableRow] + + val result = cs.query(query)(20.0f).compile.toVector.unsafeRunSync() + result.foreach(_.floatColumn should be >= 20.0f) + result should not be empty + } + + } + + + "ANN vector search" - { + + "ORDER BY ANN OF" in withCassandra5 { cs => + val table = + ks.table[VectorTableRow] + .partition(Symbol("intColumn")) + .indexBySAIVector(Symbol("vector8FloatColumn"), "vec_ann_idx", SimilarityFunction.COSINE) + .build("ann_table") + + (for { + _ <- cs.create(ks) + _ <- cs.create(table) + } yield ()).unsafeRunSync() + + val insert = table.insert.all.build.from[VectorTableRow] + val entries = (1 to 5).map { i => + VectorTableRow.instance.copy( + intColumn = i + , longColumn = i.toLong + , vector8FloatColumn = tag[VectorSize8](Vector.fill(8)(i.toFloat / 10.0f)) + ) + } + emits(entries).flatMap(e => eval(cs.execute(insert)(e)).drain).compile.drain.unsafeRunSync() + + val query = table.query + .column(Symbol("intColumn")) + .column(Symbol("longColumn")) + .orderByAnn(Symbol("vector8FloatColumn")) + .limit(3) + .build + .fromA[Vector[Float] @@ VectorSize8] + .asTuple + + val queryVec = tag[VectorSize8](Vector.fill(8)(0.3f)) + val result = cs.query(query)(queryVec).compile.toVector.unsafeRunSync() + + result.size shouldBe 3 + } + + "ANN with partition key" in withCassandra5 { cs => + val table = + ks.table[VectorTableRow] + .partition(Symbol("intColumn")) + .indexBySAIVector(Symbol("vector8FloatColumn"), "vec_pk_ann_idx", SimilarityFunction.COSINE) + .build("ann_pk_table") + + (for { + _ <- cs.create(ks) + _ <- cs.create(table) + } yield ()).unsafeRunSync() + + val insert = table.insert.all.build.from[VectorTableRow] + val entries = (1 to 5).map { i => + VectorTableRow.instance.copy( + intColumn = 1 + , longColumn = i.toLong + , vector8FloatColumn = tag[VectorSize8](Vector.fill(8)(i.toFloat / 10.0f)) + ) + } + emits(entries).flatMap(e => eval(cs.execute(insert)(e)).drain).compile.drain.unsafeRunSync() + + val query = table.query + .column(Symbol("intColumn")) + .column(Symbol("longColumn")) + .partition + .orderByAnn(Symbol("vector8FloatColumn")) + .limit(3) + .build + .fromHList + .fromTuple[(Int, Vector[Float] @@ VectorSize8)] + .asTuple + + val queryVec = tag[VectorSize8](Vector.fill(8)(0.3f)) + val result = cs.query(query)((1, queryVec)).compile.toVector.unsafeRunSync() + + result should not be empty + result.size should be <= 3 + } + + "similarity_cosine function with ANN" in withCassandra5 { cs => + val table = + ks.table[VectorTableRow] + .partition(Symbol("intColumn")) + .indexBySAIVector(Symbol("vector8FloatColumn"), "vec_sim_cos_idx", SimilarityFunction.COSINE) + .build("sim_cos_table") + + (for { + _ <- cs.create(ks) + _ <- cs.create(table) + } yield ()).unsafeRunSync() + + val insert = table.insert.all.build.from[VectorTableRow] + val entries = (1 to 3).map { i => + VectorTableRow.instance.copy( + intColumn = i + , longColumn = i.toLong + , vector8FloatColumn = tag[VectorSize8](Vector.fill(8)(i.toFloat / 10.0f)) + ) + } + emits(entries).flatMap(e => eval(cs.execute(insert)(e)).drain).compile.drain.unsafeRunSync() + + val query = table.query + .column(Symbol("intColumn")) + .orderByAnn(Symbol("vector8FloatColumn")) + .function2AtShared(functions.similarityCosine[Float, VectorSize8], Symbol("vector8FloatColumn"), Symbol("vector8FloatColumn"), Symbol("score")) + .limit(3) + .build + .fromA[Vector[Float] @@ VectorSize8] + .asTuple + + val queryVec = tag[VectorSize8](Vector.fill(8)(0.1f)) + val result = cs.query(query)(queryVec).compile.toVector.unsafeRunSync() + + result.size shouldBe 3 + result.foreach { case (score, _) => score should (be >= 0.0f and be <= 1.0f) } + } + + "similarity_euclidean function with ANN" in withCassandra5 { cs => + val table = + ks.table[VectorTableRow] + .partition(Symbol("intColumn")) + .indexBySAIVector(Symbol("vector8FloatColumn"), "vec_sim_euc_idx", SimilarityFunction.EUCLIDEAN) + .build("sim_euc_table") + + (for { + _ <- cs.create(ks) + _ <- cs.create(table) + } yield ()).unsafeRunSync() + + val insert = table.insert.all.build.from[VectorTableRow] + val entries = (1 to 3).map { i => + VectorTableRow.instance.copy( + intColumn = i + , longColumn = i.toLong + , vector8FloatColumn = tag[VectorSize8](Vector.fill(8)(i.toFloat / 10.0f)) + ) + } + emits(entries).flatMap(e => eval(cs.execute(insert)(e)).drain).compile.drain.unsafeRunSync() + + val query = table.query + .column(Symbol("intColumn")) + .orderByAnn(Symbol("vector8FloatColumn")) + .function2AtShared(functions.similarityEuclidean[Float, VectorSize8], Symbol("vector8FloatColumn"), Symbol("vector8FloatColumn"), Symbol("score")) + .limit(3) + .build + .fromA[Vector[Float] @@ VectorSize8] + .asTuple + + val queryVec = tag[VectorSize8](Vector.fill(8)(0.1f)) + val result = cs.query(query)(queryVec).compile.toVector.unsafeRunSync() + + result.size shouldBe 3 + result.foreach { case (score, _) => score should (be >= 0.0f and be <= 1.0f) } + } + + "similarity_dot_product function with ANN" in withCassandra5 { cs => + val table = + ks.table[VectorTableRow] + .partition(Symbol("intColumn")) + .indexBySAIVector(Symbol("vector8FloatColumn"), "vec_sim_dot_idx", SimilarityFunction.DOT_PRODUCT) + .build("sim_dot_table") + + (for { + _ <- cs.create(ks) + _ <- cs.create(table) + } yield ()).unsafeRunSync() + + val insert = table.insert.all.build.from[VectorTableRow] + // Use normalized vectors for dot product (magnitudes close to 1) + val entries = (1 to 3).map { i => + val mag = math.sqrt(8.0 * (i.toFloat / 10.0f) * (i.toFloat / 10.0f)).toFloat + val norm = if (mag > 0) i.toFloat / 10.0f / mag else 0f + VectorTableRow.instance.copy( + intColumn = i + , longColumn = i.toLong + , vector8FloatColumn = tag[VectorSize8](Vector.fill(8)(norm)) + ) + } + emits(entries).flatMap(e => eval(cs.execute(insert)(e)).drain).compile.drain.unsafeRunSync() + + val query = table.query + .column(Symbol("intColumn")) + .orderByAnn(Symbol("vector8FloatColumn")) + .function2AtShared(functions.similarityDotProduct[Float, VectorSize8], Symbol("vector8FloatColumn"), Symbol("vector8FloatColumn"), Symbol("score")) + .limit(3) + .build + .fromA[Vector[Float] @@ VectorSize8] + .asTuple + + val norm = (1.0f / math.sqrt(8.0)).toFloat + val queryVec = tag[VectorSize8](Vector.fill(8)(norm)) + val result = cs.query(query)(queryVec).compile.toVector.unsafeRunSync() + + result.size shouldBe 3 + result.foreach { case (score, _) => score should (be >= 0.0f and be <= 1.0f) } + } + } +} diff --git a/test/src/test/scala/spinoco/fs2/cassandra/builder/QueryBuilderSpec.scala b/test/src/test/scala/spinoco/fs2/cassandra/builder/QueryBuilderSpec.scala index a616c9c..186ece3 100644 --- a/test/src/test/scala/spinoco/fs2/cassandra/builder/QueryBuilderSpec.scala +++ b/test/src/test/scala/spinoco/fs2/cassandra/builder/QueryBuilderSpec.scala @@ -1,7 +1,8 @@ package spinoco.fs2.cassandra.builder import shapeless.LabelledGeneric -import spinoco.fs2.cassandra.sample.SimpleTableRow +import spinoco.fs2.cassandra.sample.{ListTableRow, SimpleTableRow, VectorTableRow} +import spinoco.fs2.cassandra.sample.VectorSizes._ import spinoco.fs2.cassandra.support.Fs2CassandraSpec import spinoco.fs2.cassandra.{Comparison, KeySpace, functions} @@ -249,4 +250,141 @@ class QueryBuilderSpec extends Fs2CassandraSpec { } + "SAI query features" - { + + val vectorTable = + ks.table[VectorTableRow] + .partition(Symbol("intColumn")) + .cluster(Symbol("longColumn")) + .indexBySAIVector(Symbol("vector8FloatColumn"), "vector_sai_idx") + .build("test_table") + + val indexedSimpleTable = + ks.table[SimpleTableRow] + .partition(Symbol("intColumn")) + .cluster(Symbol("longColumn")) + .indexBySAI(Symbol("stringColumn"), "string_idx") + .build("test_table") + + val indexedListTable = + ks.table[ListTableRow] + .partition(Symbol("intColumn")) + .cluster(Symbol("longColumn")) + .indexBySAICollection(Symbol("listColumn"), "list_idx", CollectionIndexTarget.Values) + .build("test_table") + + "will select with ANN ORDER BY" in { + vectorTable.query + .all + .partition + .orderByAnn(Symbol("vector8FloatColumn")) + .limit(10) + .build + .cqlStatement shouldBe + "SELECT intColumn,longColumn,vector4IntColumn,vector8FloatColumn FROM test_ks.test_table" + + " WHERE intColumn = :intColumn ORDER BY vector8FloatColumn ANN OF :vector8FloatColumn LIMIT 10" + } + + "will select with ANN ORDER BY aliased" in { + vectorTable.query + .all + .partition + .orderByAnn(Symbol("vector8FloatColumn"), Symbol("queryVec")) + .limit(10) + .build + .cqlStatement shouldBe + "SELECT intColumn,longColumn,vector4IntColumn,vector8FloatColumn FROM test_ks.test_table" + + " WHERE intColumn = :intColumn ORDER BY vector8FloatColumn ANN OF :queryVec LIMIT 10" + } + + "will select with similarity function" in { + vectorTable.query + .function2At(functions.similarityCosine[Float, VectorSize8], Symbol("vector8FloatColumn"), Symbol("queryVec"), Symbol("score")) + .partition + .build + .cqlStatement shouldBe + "SELECT similarity_cosine(vector8FloatColumn, :queryVec) AS score FROM test_ks.test_table" + + " WHERE intColumn = :intColumn" + } + + "will select with ANN and shared similarity function" in { + vectorTable.query + .all + .partition + .orderByAnn(Symbol("vector8FloatColumn")) + .function2AtShared(functions.similarityCosine[Float, VectorSize8], Symbol("vector8FloatColumn"), Symbol("vector8FloatColumn"), Symbol("score")) + .limit(10) + .build + .cqlStatement shouldBe + "SELECT intColumn,longColumn,vector4IntColumn,vector8FloatColumn,similarity_cosine(vector8FloatColumn, :vector8FloatColumn) AS score FROM test_ks.test_table" + + " WHERE intColumn = :intColumn ORDER BY vector8FloatColumn ANN OF :vector8FloatColumn LIMIT 10" + } + + "will select with similarity euclidean" in { + vectorTable.query + .function2At(functions.similarityEuclidean[Float, VectorSize8], Symbol("vector8FloatColumn"), Symbol("queryVec"), Symbol("score")) + .partition + .build + .cqlStatement shouldBe + "SELECT similarity_euclidean(vector8FloatColumn, :queryVec) AS score FROM test_ks.test_table" + + " WHERE intColumn = :intColumn" + } + + "will select with similarity dot product" in { + vectorTable.query + .function2At(functions.similarityDotProduct[Float, VectorSize8], Symbol("vector8FloatColumn"), Symbol("queryVec"), Symbol("score")) + .partition + .build + .cqlStatement shouldBe + "SELECT similarity_dot_product(vector8FloatColumn, :queryVec) AS score FROM test_ks.test_table" + + " WHERE intColumn = :intColumn" + } + + "will select with byIndex CONTAINS" in { + indexedListTable.query + .all + .partition + .byIndex(Symbol("listColumn"), Comparison.CONTAINS) + .build + .cqlStatement shouldBe + "SELECT intColumn,longColumn,listColumn,setColumn,vectorColumn,seqColumn FROM test_ks.test_table" + + " WHERE intColumn = :intColumn AND listColumn CONTAINS :listColumn" + } + + "will select with byIndex EQ on SAI" in { + indexedSimpleTable.query + .all + .partition + .byIndex(Symbol("stringColumn"), Comparison.EQ) + .build + .cqlStatement shouldBe + "SELECT intColumn,longColumn,stringColumn,asciiColumn,floatColumn,doubleColumn,bigDecimalColumn,bigIntColumn,blobColumn,uuidColumn,timeUuidColumn,durationColumn,inetAddressColumn,enumColumn FROM test_ks.test_table" + + " WHERE intColumn = :intColumn AND stringColumn = :stringColumn" + } + + "will select with byIndexIn" in { + indexedSimpleTable.query + .all + .partition + .byIndexIn(Symbol("stringColumn")) + .build + .cqlStatement shouldBe + "SELECT intColumn,longColumn,stringColumn,asciiColumn,floatColumn,doubleColumn,bigDecimalColumn,bigIntColumn,blobColumn,uuidColumn,timeUuidColumn,durationColumn,inetAddressColumn,enumColumn FROM test_ks.test_table" + + " WHERE intColumn = :intColumn AND stringColumn IN :stringColumn" + } + + "will select with byIndexIn aliased" in { + indexedSimpleTable.query + .all + .partition + .byIndexIn(Symbol("stringColumn"), Symbol("statusList")) + .build + .cqlStatement shouldBe + "SELECT intColumn,longColumn,stringColumn,asciiColumn,floatColumn,doubleColumn,bigDecimalColumn,bigIntColumn,blobColumn,uuidColumn,timeUuidColumn,durationColumn,inetAddressColumn,enumColumn FROM test_ks.test_table" + + " WHERE intColumn = :intColumn AND stringColumn IN :statusList" + } + + } + + } diff --git a/test/src/test/scala/spinoco/fs2/cassandra/builder/TableBuilderSpec.scala b/test/src/test/scala/spinoco/fs2/cassandra/builder/TableBuilderSpec.scala index 990f0da..17db8e1 100644 --- a/test/src/test/scala/spinoco/fs2/cassandra/builder/TableBuilderSpec.scala +++ b/test/src/test/scala/spinoco/fs2/cassandra/builder/TableBuilderSpec.scala @@ -212,5 +212,111 @@ class TableBuilderSpec extends Fs2CassandraSpec{ } + "DDL for table with SAI indexes with " - { + + val simpleTableDef = "CREATE TABLE test_ks.test_table (intColumn int,longColumn bigint,stringColumn text,asciiColumn ascii,floatColumn float,doubleColumn double,bigDecimalColumn decimal,bigIntColumn varint,blobColumn blob,uuidColumn uuid,timeUuidColumn timeuuid,durationColumn bigint,inetAddressColumn inet,enumColumn text," + + "SAI index" in { + val table = + ks.table[SimpleTableRow] + .partition(Symbol("intColumn")) + .indexBySAI(Symbol("asciiColumn"), "ascii_sai_idx") + .build("test_table") + + table.cqlStatement.toSet shouldBe Set( + s"$simpleTableDef PRIMARY KEY ((intColumn)))" + , "CREATE CUSTOM INDEX ascii_sai_idx ON test_ks.test_table (asciiColumn) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex'" + ) + } + + "SAI index with options" in { + val table = + ks.table[SimpleTableRow] + .partition(Symbol("intColumn")) + .indexBySAI(Symbol("stringColumn"), "string_sai_idx", Map("case_sensitive" -> "false", "normalize" -> "true")) + .build("test_table") + + table.cqlStatement.toSet shouldBe Set( + s"$simpleTableDef PRIMARY KEY ((intColumn)))" + , "CREATE CUSTOM INDEX string_sai_idx ON test_ks.test_table (stringColumn) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex' WITH OPTIONS = {'case_sensitive': 'false','normalize': 'true'}" + ) + } + + "SAI vector index with cosine similarity" in { + val vectorTableDef = "CREATE TABLE test_ks.test_table (intColumn int,longColumn bigint,vector4IntColumn vector,vector8FloatColumn vector," + + val table = + ks.table[VectorTableRow] + .partition(Symbol("intColumn")) + .indexBySAIVector(Symbol("vector8FloatColumn"), "vector_sai_idx", SimilarityFunction.COSINE) + .build("test_table") + + table.cqlStatement.toSet shouldBe Set( + s"$vectorTableDef PRIMARY KEY ((intColumn)))" + , "CREATE CUSTOM INDEX vector_sai_idx ON test_ks.test_table (vector8FloatColumn) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex' WITH OPTIONS = {'similarity_function': 'COSINE'}" + ) + } + + "SAI vector index with dot product similarity" in { + val vectorTableDef = "CREATE TABLE test_ks.test_table (intColumn int,longColumn bigint,vector4IntColumn vector,vector8FloatColumn vector," + + val table = + ks.table[VectorTableRow] + .partition(Symbol("intColumn")) + .indexBySAIVector(Symbol("vector8FloatColumn"), "vector_sai_idx", SimilarityFunction.DOT_PRODUCT) + .build("test_table") + + table.cqlStatement.toSet shouldBe Set( + s"$vectorTableDef PRIMARY KEY ((intColumn)))" + , "CREATE CUSTOM INDEX vector_sai_idx ON test_ks.test_table (vector8FloatColumn) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex' WITH OPTIONS = {'similarity_function': 'DOT_PRODUCT'}" + ) + } + + "SAI collection index with KEYS" in { + val mapTableDef = "CREATE TABLE test_ks.test_table (intColumn int,longColumn bigint,mapStringColumn map,mapIntColumn map," + + val table = + ks.table[MapTableRow] + .partition(Symbol("intColumn")) + .indexBySAICollection(Symbol("mapStringColumn"), "map_keys_idx", CollectionIndexTarget.Keys) + .build("test_table") + + table.cqlStatement.toSet shouldBe Set( + s"$mapTableDef PRIMARY KEY ((intColumn)))" + , "CREATE CUSTOM INDEX map_keys_idx ON test_ks.test_table (KEYS(mapStringColumn)) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex'" + ) + } + + "SAI collection index with VALUES" in { + val listTableDef = "CREATE TABLE test_ks.test_table (intColumn int,longColumn bigint,listColumn list,setColumn set,vectorColumn list,seqColumn list," + + val table = + ks.table[ListTableRow] + .partition(Symbol("intColumn")) + .indexBySAICollection(Symbol("listColumn"), "list_values_idx", CollectionIndexTarget.Values) + .build("test_table") + + table.cqlStatement.toSet shouldBe Set( + s"$listTableDef PRIMARY KEY ((intColumn)))" + , "CREATE CUSTOM INDEX list_values_idx ON test_ks.test_table (VALUES(listColumn)) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex'" + ) + } + + "SAI collection index with ENTRIES" in { + val mapTableDef = "CREATE TABLE test_ks.test_table (intColumn int,longColumn bigint,mapStringColumn map,mapIntColumn map," + + val table = + ks.table[MapTableRow] + .partition(Symbol("intColumn")) + .indexBySAICollection(Symbol("mapStringColumn"), "map_entries_idx", CollectionIndexTarget.Entries) + .build("test_table") + + table.cqlStatement.toSet shouldBe Set( + s"$mapTableDef PRIMARY KEY ((intColumn)))" + , "CREATE CUSTOM INDEX map_entries_idx ON test_ks.test_table (ENTRIES(mapStringColumn)) USING 'org.apache.cassandra.index.sai.StorageAttachedIndex'" + ) + } + } + }