Further improvements to Materialized support

outr · Jun 9, 2024 · 0b225da · 0b225da
1 parent 4d7e84e
commit 0b225da
Show file tree

Hide file tree

Showing 7 changed files with 15 additions and 12 deletions.
diff --git a/all/src/test/scala/spec/SimpleHaloAndSQLiteSpec.scala b/all/src/test/scala/spec/SimpleHaloAndSQLiteSpec.scala
@@ -179,7 +179,7 @@ class SimpleHaloAndSQLiteSpec extends AsyncWordSpec with AsyncIOSpec with Matche
     }
     "query for materialized indexes" in {
       Person.withSearchContext { implicit context =>
-        Person.query.materialized.compile.toList.map { list =>
+        Person.query.materialized(Person.age).compile.toList.map { list =>
           list.map(m => m(Person.age)).toSet should be(Set(19, 21))
         }
       }

diff --git a/core/src/main/scala/lightdb/index/Index.scala b/core/src/main/scala/lightdb/index/Index.scala
@@ -13,7 +13,6 @@ trait Index[F, D <: Document[D]] {
 
   def fieldName: String
   def indexSupport: IndexSupport[D]
-  def materialize: Boolean
   def get: D => List[F]
   def getJson: D => List[Json] = (doc: D) => get(doc).map(_.json)
 

diff --git a/core/src/main/scala/lightdb/query/Query.scala b/core/src/main/scala/lightdb/query/Query.scala
@@ -18,6 +18,7 @@ case class Query[D <: Document[D], V](indexSupport: IndexSupport[D],
                                       offset: Int = 0,
                                       pageSize: Int = 1_000,
                                       limit: Option[Int] = None,
+                                      materializedIndexes: List[Index[_, D]] = Nil,
                                       countTotal: Boolean = true) {
   def evalConvert[T](converter: D => IO[T]): Query[D, T] = copy(convert = converter)
   def convert[T](converter: D => T): Query[D, T] = copy(convert = doc => IO.blocking(converter(doc)))
@@ -98,7 +99,10 @@ case class Query[D <: Document[D], V](indexSupport: IndexSupport[D],
 
   def idStream(implicit context: SearchContext[D]): fs2.Stream[IO, Id[D]] = pageStream.flatMap(_.idStream)
 
-  def materialized(implicit context: SearchContext[D]): fs2.Stream[IO, Materialized[D]] = pageStream.flatMap(_.materializedStream)
+  def materialized(indexes: Index[_, D]*)
+                  (implicit context: SearchContext[D]): fs2.Stream[IO, Materialized[D]] = {
+    copy(materializedIndexes = indexes.toList).pageStream.flatMap(_.materializedStream)
+  }
 
   def stream(implicit context: SearchContext[D]): fs2.Stream[IO, V] = pageStream.flatMap(_.stream)
 

diff --git a/lucene/src/main/scala/lightdb/lucene/LuceneIndex.scala b/lucene/src/main/scala/lightdb/lucene/LuceneIndex.scala
@@ -25,8 +25,6 @@ case class LuceneIndex[F, D <: Document[D]](fieldName: String,
 
   private implicit def filter2Lucene(filter: Filter[D]): LuceneFilter[D] = filter.asInstanceOf[LuceneFilter[D]]
 
-  override def materialize: Boolean = false   // TODO: Support materialization
-
   lazy val fieldSortName: String = {
     val separate = rw.definition.className.collect {
       case "lightdb.spatial.GeoPoint" => true

diff --git a/sql/src/main/scala/lightdb/sql/SQLIndex.scala b/sql/src/main/scala/lightdb/sql/SQLIndex.scala
@@ -7,7 +7,6 @@ import lightdb.query.Filter
 
 case class SQLIndex[F, D <: Document[D]](fieldName: String,
                                          indexSupport: IndexSupport[D],
-                                         materialize: Boolean,
                                          get: D => List[F])(implicit val rw: RW[F]) extends Index[F, D] {
   override def is(value: F): Filter[D] = SQLFilter[D](s"$fieldName = ?", List(value.json))
 

diff --git a/sql/src/main/scala/lightdb/sql/SQLIndexer.scala b/sql/src/main/scala/lightdb/sql/SQLIndexer.scala
@@ -19,7 +19,6 @@ case class SQLIndexer[D <: Document[D]](indexSupport: SQLSupport[D]) extends Ind
               (implicit rw: RW[F]): Index[F, D] = SQLIndex(
     fieldName = name,
     indexSupport = indexSupport,
-    materialize = materialize,
     get = doc => get(doc)
   )
 

diff --git a/sql/src/main/scala/lightdb/sql/SQLSupport.scala b/sql/src/main/scala/lightdb/sql/SQLSupport.scala
@@ -140,10 +140,14 @@ trait SQLSupport[D <: Document[D]] extends IndexSupport[D] {
       case Nil => ""
       case list => list.mkString("ORDER BY ", ", ", "")
     }
-    val fields = index.fields.filter(_.materialize).map(_.fieldName).mkString(", ")
+    val indexes = query.materializedIndexes match {
+      case l if !l.contains(_id) => _id :: l
+      case l => l
+    }
+    val fieldNames = indexes.map(_.fieldName).mkString(", ")
     val sql =
       s"""SELECT
-         |  $fields
+         |  $fieldNames
          |FROM
          |  ${collection.collectionName}
          |$filters
@@ -156,7 +160,7 @@ trait SQLSupport[D <: Document[D]] extends IndexSupport[D] {
     val ps = prepare(sql, params)
     val rs = ps.executeQuery()
     try {
-      val materialized = this.materialized(rs)
+      val materialized = this.materialized(rs, indexes)
       PagedResults(
         query = query,
         context = SQLPageContext(context),
@@ -172,12 +176,12 @@ trait SQLSupport[D <: Document[D]] extends IndexSupport[D] {
     }
   }
 
-  protected def materialized(rs: ResultSet): List[Materialized[D]] = {
+  protected def materialized(rs: ResultSet, indexes: List[Index[_, D]]): List[Materialized[D]] = {
     val iterator = new Iterator[Materialized[D]] {
       override def hasNext: Boolean = rs.next()
 
       override def next(): Materialized[D] = {
-        val map = index.fields.filter(_.materialize).map { index =>
+        val map = indexes.map { index =>
           index.fieldName -> getJson(rs, index.fieldName)
         }.toMap
         Materialized[D](Obj(map))