Lots of updates

outr · May 28, 2024 · 133a604 · 133a604
1 parent 5df094d
commit 133a604
Show file tree

Hide file tree

Showing 12 changed files with 79 additions and 15 deletions.
diff --git a/all/src/test/scala/spec/SimpleHaloAndLuceneSpec.scala b/all/src/test/scala/spec/SimpleHaloAndLuceneSpec.scala
@@ -206,7 +206,7 @@ class SimpleHaloAndLuceneSpec extends AsyncWordSpec with AsyncIOSpec with Matche
     }
     "do paginated search as a stream converting to name" in {
       Person.withSearchContext { implicit context =>
-        Person.query.convert(p => IO.pure(p.name)).pageSize(1).countTotal(true).stream.compile.toList.map { names =>
+        Person.query.convert(_.name).pageSize(1).countTotal(true).stream.compile.toList.map { names =>
           names.length should be(3)
           names.toSet should be(Set("John Doe", "Jane Doe", "Bob Dole"))
         }
@@ -235,7 +235,7 @@ class SimpleHaloAndLuceneSpec extends AsyncWordSpec with AsyncIOSpec with Matche
     }
     "search using tokenized data and a parsed query" in {
       Person.query
-        .filter(Person.search parsed "joh% 21")
+        .filter(Person.search.words("joh 21"))
         .toList
         .map { results =>
           results.map(_.name) should be(List("John Doe"))

diff --git a/core/src/main/scala/lightdb/index/IndexSupport.scala b/core/src/main/scala/lightdb/index/IndexSupport.scala
@@ -39,6 +39,7 @@ trait IndexSupport[D <: Document[D]] extends DocumentModel[D] {
   def doSearch[V](query: Query[D, V],
                   context: SearchContext[D],
                   offset: Int,
+                  limit: Option[Int],
                   after: Option[PagedResults[D, V]]): IO[PagedResults[D, V]]
 
   protected def indexDoc(doc: D, fields: List[IndexedField[_, D]]): IO[Unit]

diff --git a/core/src/main/scala/lightdb/model/AbstractCollection.scala b/core/src/main/scala/lightdb/model/AbstractCollection.scala
@@ -60,6 +60,8 @@ trait AbstractCollection[D <: Document[D]] extends DocumentActionSupport[D] {
     }
   }
 
+  def set(docs: Seq[D]): IO[Int] = docs.map(set).sequence.map(_.size)
+
   def modify(id: Id[D])
             (f: Option[D] => IO[Option[D]])
             (implicit existingLock: DocLock[D] = new DocLock.Empty[D]): IO[Option[D]] = withLock(id) { implicit lock =>

diff --git a/core/src/main/scala/lightdb/query/PageContext.scala b/core/src/main/scala/lightdb/query/PageContext.scala
@@ -11,6 +11,7 @@ trait PageContext[D <: Document[D]] {
       query = currentPage.query,
       context = context,
       offset = currentPage.offset + currentPage.query.pageSize,
+      limit = currentPage.query.limit,
       after = Some(currentPage)
     ).map(Some.apply)
   } else {

diff --git a/core/src/main/scala/lightdb/query/PagedResults.scala b/core/src/main/scala/lightdb/query/PagedResults.scala
@@ -11,7 +11,7 @@ case class PagedResults[D <: Document[D], V](query: Query[D, V],
                                              idsAndScores: List[(Id[D], Double)],
                                              getter: Option[Id[D] => IO[D]] = None) {
   lazy val page: Int = offset / query.pageSize
-  lazy val pages: Int = math.ceil(total.toDouble / query.pageSize.toDouble).toInt
+  lazy val pages: Int = math.ceil(query.limit.getOrElse(total).toDouble / query.pageSize.toDouble).toInt
 
   lazy val ids: List[Id[D]] = idsAndScores.map(_._1)
   lazy val scores: List[Double] = idsAndScores.map(_._2)

diff --git a/core/src/main/scala/lightdb/query/Query.scala b/core/src/main/scala/lightdb/query/Query.scala
@@ -16,8 +16,10 @@ case class Query[D <: Document[D], V](indexSupport: IndexSupport[D],
                                       scoreDocs: Boolean = false,
                                       offset: Int = 0,
                                       pageSize: Int = 1_000,
+                                      limit: Option[Int] = None,
                                       countTotal: Boolean = true) {
-  def convert[T](converter: D => IO[T]): Query[D, T] = copy(convert = converter)
+  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(converter(doc)))
 
   def filter(filter: Filter[D], and: Boolean = false): Query[D, V] = {
     if (and && this.filter.nonEmpty) {
@@ -27,13 +29,23 @@ case class Query[D <: Document[D], V](indexSupport: IndexSupport[D],
     }
   }
 
+  def filters(filters: Filter[D]*): Query[D, V] = if (filters.nonEmpty) {
+    var filter = filters.head
+    filters.tail.foreach { f =>
+      filter = filter && f
+    }
+    this.filter(filter)
+  } else {
+    this
+  }
+
   def sort(sort: Sort*): Query[D, V] = copy(sort = this.sort ::: sort.toList)
 
   def distance(field: IndexedField[GeoPoint, D],
                from: GeoPoint,
                sort: Boolean = true,
                radius: Option[Length] = None): Query[D, DistanceAndDoc[D]] = {
-    var q = convert(doc => IO {
+    var q = convert(doc => {
       DistanceAndDoc(
         doc = doc,
         distance = DistanceCalculator(from, field.get(doc).head)
@@ -60,12 +72,15 @@ case class Query[D <: Document[D], V](indexSupport: IndexSupport[D],
 
   def pageSize(size: Int): Query[D, V] = copy(pageSize = size)
 
+  def limit(limit: Int): Query[D, V] = copy(limit = Some(limit))
+
   def countTotal(b: Boolean): Query[D, V] = copy(countTotal = b)
 
   def search()(implicit context: SearchContext[D]): IO[PagedResults[D, V]] = indexSupport.doSearch(
     query = this,
     context = context,
     offset = offset,
+    limit = limit,
     after = None
   )
 
@@ -98,6 +113,12 @@ case class Query[D <: Document[D], V](indexSupport: IndexSupport[D],
     stream.compile.toList
   }
 
+  def first: IO[Option[V]] = indexSupport.withSearchContext { implicit context =>
+    stream.take(1).compile.last
+  }
+
+  def one: IO[V] = first.map(_.getOrElse(throw new RuntimeException(s"No results for query: $this")))
+
   def count: IO[Int] = indexSupport.withSearchContext { implicit context =>
     idStream.compile.count.map(_.toInt)
   }

diff --git a/core/src/main/scala/lightdb/query/Sort.scala b/core/src/main/scala/lightdb/query/Sort.scala
@@ -9,7 +9,13 @@ trait Sort
 object Sort {
   case object BestMatch extends Sort
   case object IndexOrder extends Sort
-  case class ByField[D <: Document[D], F](field: IndexedField[F, D], reverse: Boolean = false) extends Sort
+  case class ByField[D <: Document[D], F](field: IndexedField[F, D], direction: SortDirection = SortDirection.Ascending) extends Sort {
+    def direction(direction: SortDirection): ByField[D, F] = copy(direction = direction)
+    def ascending: ByField[D, F] = direction(SortDirection.Ascending)
+    def asc: ByField[D, F] = direction(SortDirection.Ascending)
+    def descending: ByField[D, F] = direction(SortDirection.Descending)
+    def desc: ByField[D, F] = direction(SortDirection.Descending)
+  }
   case class ByDistance[D <: Document[D]](field: IndexedField[GeoPoint, D],
                                           from: GeoPoint) extends Sort
 }
diff --git a/core/src/main/scala/lightdb/query/SortDirection.scala b/core/src/main/scala/lightdb/query/SortDirection.scala
@@ -0,0 +1,8 @@
+package lightdb.query
+
+sealed trait SortDirection
+
+object SortDirection {
+  case object Ascending extends SortDirection
+  case object Descending extends SortDirection
+}
diff --git a/lucene/src/main/scala/lightdb/lucene/LuceneIndex.scala b/lucene/src/main/scala/lightdb/lucene/LuceneIndex.scala
@@ -41,10 +41,28 @@ case class LuceneIndex[F, D <: Document[D]](fieldName: String,
     LuceneFilter(() => {
       val parser = new QueryParser(fieldName, lucene.index.analyzer)
       parser.setAllowLeadingWildcard(allowLeadingWildcard)
+      parser.setSplitOnWhitespace(true)
       parser.parse(query)
     })
   }
 
+  def words(s: String,
+            matchStartsWith: Boolean = true,
+            matchEndsWith: Boolean = false): LuceneFilter[D] = {
+    val words = s.split("\\s+").map { w =>
+      if (matchStartsWith && matchEndsWith) {
+        s"%$w%"
+      } else if (matchStartsWith) {
+        s"%$w"
+      } else if (matchEndsWith) {
+        s"$w%"
+      } else {
+        w
+      }
+    }.mkString(" ")
+    parsed(words, allowLeadingWildcard = matchEndsWith)
+  }
+
   def IN(values: Seq[F]): LuceneFilter[D] = {
     val b = new BooleanQuery.Builder
     b.setMinimumNumberShouldMatch(1)

diff --git a/lucene/src/main/scala/lightdb/lucene/LuceneSupport.scala b/lucene/src/main/scala/lightdb/lucene/LuceneSupport.scala
@@ -4,7 +4,7 @@ import cats.effect.IO
 import fabric.define.DefType
 import lightdb._
 import lightdb.index.{IndexSupport, IndexedField}
-import lightdb.query.{Filter, PageContext, PagedResults, Query, SearchContext, Sort}
+import lightdb.query.{Filter, PageContext, PagedResults, Query, SearchContext, Sort, SortDirection}
 import lightdb.spatial.GeoPoint
 import org.apache.lucene.document.{LatLonDocValuesField, LatLonPoint}
 import org.apache.lucene.search.{IndexSearcher, MatchAllDocsQuery, ScoreDoc, SearcherFactory, SearcherManager, SortField, SortedNumericSortField, TopFieldDocs, Query => LuceneQuery, Sort => LuceneSort}
@@ -21,11 +21,11 @@ trait LuceneSupport[D <: Document[D]] extends IndexSupport[D] {
   private def sort2SortField(sort: Sort): SortField = sort match {
     case Sort.BestMatch => SortField.FIELD_SCORE
     case Sort.IndexOrder => SortField.FIELD_DOC
-    case Sort.ByField(field, reverse) =>
+    case Sort.ByField(field, dir) =>
       val f = field.asInstanceOf[LuceneIndex[_, D]]
       f.rw.definition match {
-        case DefType.Int => new SortedNumericSortField(field.fieldName, f.sortType, reverse)
-        case DefType.Str => new SortField(field.fieldName, f.sortType, reverse)
+        case DefType.Int => new SortedNumericSortField(field.fieldName, f.sortType, dir == SortDirection.Descending)
+        case DefType.Str => new SortField(field.fieldName, f.sortType, dir == SortDirection.Descending)
         case d => throw new RuntimeException(s"Unsupported sort definition: $d")
       }
     case Sort.ByDistance(field, from) =>
@@ -36,6 +36,7 @@ trait LuceneSupport[D <: Document[D]] extends IndexSupport[D] {
   override def doSearch[V](query: Query[D, V],
                            context: SearchContext[D],
                            offset: Int,
+                           limit: Option[Int],
                            after: Option[PagedResults[D, V]]): IO[PagedResults[D, V]] = IO {
     val q = query.filter.map(_.asInstanceOf[LuceneFilter[D]].asQuery()).getOrElse(new MatchAllDocsQuery)
     val sortFields = query.sort match {
@@ -50,7 +51,7 @@ trait LuceneSupport[D <: Document[D]] extends IndexSupport[D] {
         indexSearcher.searchAfter(afterDoc, q, query.pageSize, s, query.scoreDocs)
       case None => indexSearcher.search(q, query.pageSize, s, query.scoreDocs)
     }
-    val scoreDocs: List[ScoreDoc] = topFieldDocs.scoreDocs.toList
+    val scoreDocs: List[ScoreDoc] = topFieldDocs.scoreDocs.toList.take(limit.getOrElse(Int.MaxValue) - offset)
     val total: Int = topFieldDocs.totalHits.value.toInt
     val storedFields: StoredFields = indexSearcher.storedFields()
     val idsAndScores = scoreDocs.map(doc => Id[D](storedFields.document(doc.doc).get("_id")) -> doc.score.toDouble)

diff --git a/sqlite/src/main/scala/lightdb/sqlite/SQLIndexedField.scala b/sqlite/src/main/scala/lightdb/sqlite/SQLIndexedField.scala
@@ -11,6 +11,11 @@ case class SQLIndexedField[F, D <: Document[D]](fieldName: String,
 
   def is(value: F): SQLFilter[D] = SQLFilter[D](s"$fieldName = ?", List(value.json))
 
+  def >(value: F): SQLFilter[D] = SQLFilter[D](s"$fieldName > ?", List(value.json))
+  def >=(value: F): SQLFilter[D] = SQLFilter[D](s"$fieldName >= ?", List(value.json))
+  def <(value: F): SQLFilter[D] = SQLFilter[D](s"$fieldName < ?", List(value.json))
+  def <=(value: F): SQLFilter[D] = SQLFilter[D](s"$fieldName <= ?", List(value.json))
+
   def between(v1: F, v2: F): SQLFilter[D] = SQLFilter[D](s"$fieldName BETWEEN ? AND ?", List(v1.json, v2.json))
 
   def IN(values: Seq[F]): SQLFilter[D] = SQLFilter[D](s"$fieldName IN (${values.map(_ => "?").mkString(", ")})", values.toList.map(_.json))

diff --git a/sqlite/src/main/scala/lightdb/sqlite/SQLiteSupport.scala b/sqlite/src/main/scala/lightdb/sqlite/SQLiteSupport.scala
@@ -6,7 +6,7 @@ import fabric.io.JsonFormatter
 import lightdb.{Document, Id}
 import lightdb.index.{IndexSupport, IndexedField}
 import lightdb.model.AbstractCollection
-import lightdb.query.{PagedResults, Query, SearchContext, Sort}
+import lightdb.query.{PagedResults, Query, SearchContext, Sort, SortDirection}
 import lightdb.util.FlushingBacklog
 
 import java.nio.file.{Files, Path}
@@ -96,6 +96,7 @@ trait SQLiteSupport[D <: Document[D]] extends IndexSupport[D] {
   override def doSearch[V](query: Query[D, V],
                            context: SearchContext[D],
                            offset: Int,
+                           limit: Option[Int],
                            after: Option[PagedResults[D, V]]): IO[PagedResults[D, V]] = IO {
     var params = List.empty[Json]
     val filters = query.filter match {
@@ -124,8 +125,8 @@ trait SQLiteSupport[D <: Document[D]] extends IndexSupport[D] {
       -1
     }
     val sort = query.sort.collect {
-      case Sort.ByField(field, reverse) =>
-        val dir = if (reverse) "DESC" else "ASC"
+      case Sort.ByField(field, direction) =>
+        val dir = if (direction == SortDirection.Descending) "DESC" else "ASC"
         s"${field.fieldName} $dir"
     } match {
       case Nil => ""
@@ -139,7 +140,7 @@ trait SQLiteSupport[D <: Document[D]] extends IndexSupport[D] {
          |$filters
          |$sort
          |LIMIT
-         |  ${query.pageSize}
+         |  ${query.limit.getOrElse(query.pageSize)}
          |OFFSET
          |  $offset
          |""".stripMargin