Skip to content

Commit

Permalink
Better spatial support
Browse files Browse the repository at this point in the history
  • Loading branch information
@darkfrog26
darkfrog26 committed Sep 9, 2024
1 parent 3cf081a commit 5480c68
Show file tree
Hide file tree
Showing 11 changed files with 142 additions and 89 deletions.
33 changes: 17 additions & 16 deletions async/src/main/scala/lightdb/async/AsyncQuery.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,8 @@ case class AsyncQuery[Doc <: Document[Doc], Model <: DocumentModel[Doc]](collect
limit: Option[Int] = None,
countTotal: Boolean = false,
scoreDocs: Boolean = false,
minDocScore: Option[Double] = None) { query =>
minDocScore: Option[Double] = None) {
query =>
private[async] def toQuery: Query[Doc, Model] = Query[Doc, Model](collection, filter, sort, offset, limit, countTotal, scoreDocs, minDocScore)

def scored: AsyncQuery[Doc, Model] = copy(scoreDocs = true)
Expand Down Expand Up @@ -82,11 +83,11 @@ case class AsyncQuery[Doc <: Document[Doc], Model <: DocumentModel[Doc]](collect
(implicit transaction: Transaction[Doc]): fs2.Stream[IO, (MaterializedIndex[Doc, Model], Double)] =
apply(Conversion.Materialized[Doc, Model](f(collection.model)))

def distance(f: Model => Field[Doc, Option[Geo.Point]],
from: Geo.Point,
sort: Boolean = true,
radius: Option[Distance] = None)
(implicit transaction: Transaction[Doc]): fs2.Stream[IO, (DistanceAndDoc[Doc], Double)] =
def distance[G <: Geo](f: Model => Field[Doc, List[G]],
from: Geo.Point,
sort: Boolean = true,
radius: Option[Distance] = None)
(implicit transaction: Transaction[Doc]): fs2.Stream[IO, (DistanceAndDoc[Doc], Double)] =
apply(Conversion.Distance(f(collection.model), from, sort, radius))
}

Expand Down Expand Up @@ -116,11 +117,11 @@ case class AsyncQuery[Doc <: Document[Doc], Model <: DocumentModel[Doc]](collect
(implicit transaction: Transaction[Doc]): fs2.Stream[IO, MaterializedIndex[Doc, Model]] =
apply(Conversion.Materialized[Doc, Model](f(collection.model)))

def distance(f: Model => Field[Doc, Option[Geo.Point]],
from: Geo.Point,
sort: Boolean = true,
radius: Option[Distance] = None)
(implicit transaction: Transaction[Doc]): fs2.Stream[IO, DistanceAndDoc[Doc]] =
def distance[G <: Geo](f: Model => Field[Doc, List[G]],
from: Geo.Point,
sort: Boolean = true,
radius: Option[Distance] = None)
(implicit transaction: Transaction[Doc]): fs2.Stream[IO, DistanceAndDoc[Doc]] =
apply(Conversion.Distance(f(collection.model), from, sort, radius))
}

Expand Down Expand Up @@ -159,11 +160,11 @@ case class AsyncQuery[Doc <: Document[Doc], Model <: DocumentModel[Doc]](collect
(implicit transaction: Transaction[Doc]): IO[AsyncSearchResults[Doc, MaterializedIndex[Doc, Model]]] =
apply(Conversion.Materialized(f(collection.model)))

def distance(f: Model => Field[Doc, Option[Geo.Point]],
from: Geo.Point,
sort: Boolean = true,
radius: Option[Distance] = None)
(implicit transaction: Transaction[Doc]): IO[AsyncSearchResults[Doc, DistanceAndDoc[Doc]]] =
def distance[G <: Geo](f: Model => Field[Doc, List[G]],
from: Geo.Point,
sort: Boolean = true,
radius: Option[Distance] = None)
(implicit transaction: Transaction[Doc]): IO[AsyncSearchResults[Doc, DistanceAndDoc[Doc]]] =
apply(Conversion.Distance(f(collection.model), from, sort, radius))
}

Expand Down
19 changes: 12 additions & 7 deletions core/src/main/scala/lightdb/Field.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,15 @@ sealed class Field[Doc <: Document[Doc], V](val name: String,
case _ => false
}

lazy val className: Option[String] = rw.definition match {
case DefType.Opt(DefType.Obj(_, Some(cn))) => Some(cn)
case DefType.Obj(_, Some(cn)) => Some(cn)
case DefType.Opt(DefType.Poly(_, Some(cn))) => Some(cn)
case DefType.Poly(_, Some(cn)) => Some(cn)
case _ => None
lazy val className: Option[String] = {
def lookup(d: DefType): Option[String] = d match {
case DefType.Opt(d) => lookup(d)
case DefType.Arr(d) => lookup(d)
case DefType.Poly(_, cn) => cn
case DefType.Obj(_, cn) => cn
case _ => None
}
lookup(rw.definition)
}

lazy val isSpatial: Boolean = className.exists(_.startsWith("lightdb.spatial.Geo"))
Expand Down Expand Up @@ -73,6 +76,8 @@ sealed class Field[Doc <: Document[Doc], V](val name: String,

def opt: Field[Doc, Option[V]] = new Field[Doc, Option[V]](name, doc => Option(get(doc)), () => implicitly[RW[Option[V]]], indexed)

def list: Field[Doc, List[V]] = new Field[Doc, List[V]](name, doc => List(get(doc)), () => implicitly[RW[List[V]]], indexed)

override def distance(from: Geo.Point, radius: Distance): Filter[Doc] =
Filter.Distance(name, from, radius)

Expand Down Expand Up @@ -134,7 +139,7 @@ object Field {
})
case DefType.Opt(d) => string2Json(name, s, d)
case DefType.Enum(_, _) => str(s)
case DefType.Arr(d) => arr(s.split(";;").toList.map(string2Json(name, _, d)): _*)
case DefType.Arr(d) if !s.startsWith("[") => arr(s.split(";;").toList.map(string2Json(name, _, d)): _*)
case _ => try {
JsonParser(s)
} catch {
Expand Down
9 changes: 5 additions & 4 deletions core/src/main/scala/lightdb/Query.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@ case class Query[Doc <: Document[Doc], Model <: DocumentModel[Doc]](collection:
apply(Conversion.Materialized(fields))
}

def distance[G <: Geo](f: Model => Field[Doc, Option[G]],
def distance[G <: Geo](f: Model => Field[Doc, List[G]],
from: Geo.Point,
sort: Boolean = true,
radius: Option[Distance] = None)
Expand Down Expand Up @@ -121,9 +121,10 @@ case class Query[Doc <: Document[Doc], Model <: DocumentModel[Doc]](collection:
def count(implicit transaction: Transaction[Doc]): Int = copy(limit = Some(1), countTotal = true)
.search.docs.total.get

protected def distanceSearch[G <: Geo](field: Field[Doc, Option[G]],
from: Geo.Point,
sort: Boolean, radius: Option[Distance])
protected def distanceSearch[G <: Geo](field: Field[Doc, List[G]],
from: Geo.Point,
sort: Boolean,
radius: Option[Distance])
(implicit transaction: Transaction[Doc]): SearchResults[Doc, DistanceAndDoc[Doc]] = {
search(Conversion.Distance(field, from, sort, radius))
}
Expand Down
2 changes: 1 addition & 1 deletion core/src/main/scala/lightdb/Sort.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ object Sort {
def desc: ByField[Doc, F] = direction(SortDirection.Descending)
}

case class ByDistance[Doc <: Document[Doc], G <: Geo](field: Field[Doc, Option[G]],
case class ByDistance[Doc <: Document[Doc], G <: Geo](field: Field[Doc, List[G]],
from: Geo.Point,
direction: SortDirection = SortDirection.Ascending) extends Sort {
def direction(direction: SortDirection): ByDistance[Doc, G] = copy(direction = direction)
Expand Down
2 changes: 1 addition & 1 deletion core/src/main/scala/lightdb/spatial/DistanceAndDoc.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,4 @@ package lightdb.spatial

import lightdb.distance.Distance

case class DistanceAndDoc[Doc](doc: Doc, distance: Option[Distance])
case class DistanceAndDoc[Doc](doc: Doc, distance: List[Distance])
2 changes: 1 addition & 1 deletion core/src/main/scala/lightdb/store/Conversion.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ object Conversion {

case class Converted[Doc <: Document[Doc], T](f: Doc => T) extends Conversion[Doc, T]

case class Distance[Doc <: Document[Doc], G <: Geo](field: Field[Doc, Option[G]],
case class Distance[Doc <: Document[Doc], G <: Geo](field: Field[Doc, List[G]],
from: Geo.Point,
sort: Boolean,
radius: Option[lightdb.distance.Distance]) extends Conversion[Doc, DistanceAndDoc[Doc]]
Expand Down
24 changes: 12 additions & 12 deletions core/src/test/scala/spec/AbstractSpatialSpec.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,21 +44,21 @@ abstract class AbstractSpatialSpec extends AnyWordSpec with Matchers { spec =>
name = "John Doe",
age = 21,
point = newYorkCity,
geo = None,
geo = Nil,
_id = id1
)
private val p2 = Person(
name = "Jane Doe",
age = 19,
point = noble,
geo = Some(moorePolygon),
geo = List(moorePolygon),
_id = id2
)
private val p3 = Person(
name = "Bob Dole",
age = 123,
point = yonkers,
geo = Some(chicago),
geo = List(chicago, yonkers),
_id = id3
)

Expand All @@ -79,17 +79,17 @@ abstract class AbstractSpatialSpec extends AnyWordSpec with Matchers { spec =>
"sort by distance from Oklahoma City" in {
DB.people.transaction { implicit transaction =>
val list = DB.people.query.search.distance(
_.point.opt,
_.point.list,
from = oklahomaCity,
radius = Some(1320.miles)
).iterator.toList
val people = list.map(_.doc)
val distances = list.map(_.distance.get.mi)
val distances = list.map(_.distance.map(_.mi.toInt))
people.zip(distances).map {
case (p, d) => p.name -> d
} should be(List(
"Jane Doe" -> 28.55539552714398,
"John Doe" -> 1316.1301092705082
"Jane Doe" -> List(28),
"John Doe" -> List(1316)
))
}
}
Expand All @@ -101,12 +101,12 @@ abstract class AbstractSpatialSpec extends AnyWordSpec with Matchers { spec =>
radius = Some(10_000.miles)
).iterator.toList
val people = list.map(_.doc)
val distances = list.map(_.distance.get.mi)
val distances = list.map(_.distance.map(_.mi))
people.zip(distances).map {
case (p, d) => p.name -> d
} should be(List(
"Jane Doe" -> 16.01508397712445,
"Bob Dole" -> 695.6419047674393
"Jane Doe" -> List(16.01508397712445),
"Bob Dole" -> List(695.6419047674393, 1334.038796028706)
))
}
}
Expand All @@ -133,7 +133,7 @@ abstract class AbstractSpatialSpec extends AnyWordSpec with Matchers { spec =>
case class Person(name: String,
age: Int,
point: Geo.Point,
geo: Option[Geo],
geo: List[Geo],
_id: Id[Person] = Person.id()) extends Document[Person]

object Person extends DocumentModel[Person] with JsonConversion[Person] {
Expand All @@ -142,6 +142,6 @@ abstract class AbstractSpatialSpec extends AnyWordSpec with Matchers { spec =>
val name: F[String] = field("name", _.name)
val age: F[Int] = field("age", _.age)
val point: I[Geo.Point] = field.index("point", _.point)
val geo: I[Option[Geo]] = field.index("geo", _.geo)
val geo: I[List[Geo]] = field.index("geo", _.geo)
}
}
58 changes: 36 additions & 22 deletions lucene/src/main/scala/lightdb/lucene/LuceneStore.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,17 +39,11 @@ class LuceneStore[Doc <: Document[Doc], Model <: DocumentModel[Doc]](directory:
)

override def insert(doc: Doc)(implicit transaction: Transaction[Doc]): Unit = {
val luceneFields = fields.flatMap { field =>
createLuceneFields(field, doc)
}
addDoc(id(doc), luceneFields, upsert = false)
addDoc(doc, upsert = false)
}

override def upsert(doc: Doc)(implicit transaction: Transaction[Doc]): Unit = {
val luceneFields = fields.flatMap { field =>
createLuceneFields(field, doc)
}
addDoc(id(doc), luceneFields, upsert = true)
addDoc(doc, upsert = true)
}

private def createGeoFields(field: Field[Doc, _],
Expand All @@ -71,16 +65,24 @@ class LuceneStore[Doc <: Document[Doc], Model <: DocumentModel[Doc]](directory:
val polygon = convert(p)
LatLonShape.createIndexableFields(field.name, polygon)
}
val geo = json.as[Geo]
geo match {
case p: Geo.Point => indexPoint(p)
case Geo.MultiPoint(points) => points.foreach(indexPoint)
case l: Geo.Line => indexLine(l)
case Geo.MultiLine(lines) => lines.foreach(indexLine)
case p: Geo.Polygon => indexPolygon(p)
case Geo.MultiPolygon(polygons) => polygons.foreach(indexPolygon)
val list = json match {
case Arr(value, _) => value.toList.map(_.as[Geo])
case _ => List(json.as[Geo])
}
list.foreach { geo =>
geo match {
case p: Geo.Point => indexPoint(p)
case Geo.MultiPoint(points) => points.foreach(indexPoint)
case l: Geo.Line => indexLine(l)
case Geo.MultiLine(lines) => lines.foreach(indexLine)
case p: Geo.Polygon => indexPolygon(p)
case Geo.MultiPolygon(polygons) => polygons.foreach(indexPolygon)
}
add(new LatLonPoint(field.name, geo.center.latitude, geo.center.longitude))
}
if (list.isEmpty) {
add(new LatLonPoint(field.name, 0.0, 0.0))
}
add(new LatLonPoint(field.name, geo.center.latitude, geo.center.longitude))
}
add(new StoredField(field.name, JsonFormatter.Compact(json)))
}
Expand Down Expand Up @@ -123,17 +125,26 @@ class LuceneStore[Doc <: Document[Doc], Model <: DocumentModel[Doc]](directory:
add(sorted)
case NumInt(l, _) => add(new NumericDocValuesField(fieldSortName, l))
case j if field.isSpatial && j != Null =>
val g = j.as[Geo]
add(new LatLonDocValuesField(fieldSortName, g.center.latitude, g.center.longitude))
val list = j match {
case Arr(values, _) => values.toList.map(_.as[Geo])
case _ => List(j.as[Geo])
}
list.foreach { g =>
add(new LatLonDocValuesField(fieldSortName, g.center.latitude, g.center.longitude))
}
case _ => // Ignore
}
fields
}
}

private def addDoc(id: Id[Doc], fields: List[LuceneField], upsert: Boolean): Unit = if (fields.tail.nonEmpty) {
private def addDoc(doc: Doc, upsert: Boolean): Unit = if (fields.tail.nonEmpty) {
val id = this.id(doc)
val luceneFields = fields.flatMap { field =>
createLuceneFields(field, doc)
}
val document = new LuceneDocument
fields.foreach(document.add)
luceneFields.foreach(document.add)

if (upsert) {
index.indexWriter.updateDocument(new Term("_id", id.value), document)
Expand Down Expand Up @@ -285,7 +296,10 @@ class LuceneStore[Doc <: Document[Doc], Model <: DocumentModel[Doc]](directory:
parser.setSplitOnWhitespace(true)
parser.parse(query)
case Filter.Distance(fieldName, from, radius) =>
LatLonPoint.newDistanceQuery(fieldName, from.latitude, from.longitude, radius.toMeters)
val b = new BooleanQuery.Builder
b.add(LatLonPoint.newDistanceQuery(fieldName, from.latitude, from.longitude, radius.toMeters), BooleanClause.Occur.MUST)
b.add(LatLonPoint.newBoxQuery(fieldName, 0.0, 0.0, 0.0, 0.0), BooleanClause.Occur.MUST_NOT)
b.build()
case Filter.Multi(minShould, clauses) =>
val b = new BooleanQuery.Builder
val hasShould = clauses.exists(c => c.condition == Condition.Should || c.condition == Condition.Filter)
Expand Down
2 changes: 1 addition & 1 deletion sql/src/main/scala/lightdb/sql/SQLQueryBuilder.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ case class SQLQueryBuilder[Doc <: Document[Doc]](collection: Collection[Doc, _],
SQLResults(ps.executeQuery(), combinedSql, ps)
}
} catch {
case t: Throwable => throw new SQLException(s"Error executing query: $combinedSql", t)
case t: Throwable => throw new SQLException(s"Error executing query: $combinedSql (params: ${args.mkString(" | ")})", t)
}
}
}
Expand Down
15 changes: 9 additions & 6 deletions sql/src/main/scala/lightdb/sql/SQLStore.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ import lightdb.distance.Distance
import lightdb.doc.{Document, DocumentModel, JsonConversion}
import lightdb.filter.{Condition, Filter}
import lightdb.materialized.{MaterializedAggregate, MaterializedIndex}
import lightdb.spatial.DistanceAndDoc
import lightdb.spatial.{DistanceAndDoc, Geo}
import lightdb.sql.connect.ConnectionManager
import lightdb.store.{Conversion, Store, StoreMode}
import lightdb.transaction.{Transaction, TransactionKey}
Expand Down Expand Up @@ -276,9 +276,9 @@ abstract class SQLStore[Doc <: Document[Doc], Model <: DocumentModel[Doc]] exten
jsonFromFields(fields).asInstanceOf[V]
case Conversion.Distance(field, _, _, _) =>
val fieldName = s"${field.name}Distance"
val distance = Option(rs.getObject(fieldName)).map(_.toString.toDouble).map(Distance.apply)
val distances = JsonParser(rs.getString(fieldName)).as[List[Double]].map(d => Distance(d)).toList
val doc = getDoc(rs)
DistanceAndDoc(doc, distance).asInstanceOf[V]
DistanceAndDoc(doc, distances).asInstanceOf[V]
}
}
}
Expand Down Expand Up @@ -346,9 +346,7 @@ abstract class SQLStore[Doc <: Document[Doc], Model <: DocumentModel[Doc]] exten
case Sort.ByField(index, direction) =>
val dir = if (direction == SortDirection.Descending) "DESC" else "ASC"
SQLPart(s"${index.name} $dir")
case Sort.ByDistance(field, _, direction) =>
val dir = if (direction == SortDirection.Descending) "DESC" else "ASC"
SQLPart(s"${field.name}Distance $dir")
case Sort.ByDistance(field, _, direction) => sortByDistance(field, direction)
},
limit = query.limit,
offset = query.offset
Expand All @@ -373,6 +371,11 @@ abstract class SQLStore[Doc <: Document[Doc], Model <: DocumentModel[Doc]] exten
)
}

protected def sortByDistance[G <: Geo](field: Field[_, List[G]], direction: SortDirection): SQLPart = {
val dir = if (direction == SortDirection.Descending) "DESC" else "ASC"
SQLPart(s"${field.name}Distance $dir")
}

private def aggregate2SQLQuery(query: AggregateQuery[Doc, Model])
(implicit transaction: Transaction[Doc]): SQLQueryBuilder[Doc] = {
val fields = query.functions.map { f =>
Expand Down
Loading

0 comments on commit 5480c68

Please sign in to comment.