Spatial fields and geo queries on S2 cells

Cargo.toml

@@ -91,6 +91,10 @@ datasketches = { version = "0.5.0", features = ["hll"] }
 futures-util = { version = "0.3.28", optional = true }
 futures-channel = { version = "0.3.28", optional = true }
 fnv = "1.0.7"
+num-bigint = "0.4"
+num-rational = "0.4"
+num-traits = "0.2"
+robust = "1"
 typetag = "0.2.21"
 unwrap-infallible = "1.0.0"
 serde-tuple-vec-map = "1.0.1"

examples/geo_json.rs

@@ -0,0 +1,51 @@
+use tantivy::collector::TopDocs;
+use tantivy::query::SpatialQuery;
+use tantivy::schema::{Schema, Value, SPATIAL, STORED, TEXT};
+use tantivy::{Index, IndexWriter, TantivyDocument};
+fn main() -> tantivy::Result<()> {
+    let mut schema_builder = Schema::builder();
+    schema_builder.add_json_field("properties", STORED | TEXT);
+    schema_builder.add_spatial_field("geometry", STORED | SPATIAL);
+    let schema = schema_builder.build();
+    let index = Index::create_in_ram(schema.clone());
+    let mut index_writer: IndexWriter = index.writer(50_000_000)?;
+    let doc = TantivyDocument::parse_json(
+        &schema,
+        r#"{
+            "type":"Feature",
+            "geometry":{
+                "type":"Polygon",
+                "coordinates":[[[-99.483911,45.577697],[-99.483869,45.571457],[-99.481739,45.571461],[-99.474881,45.571584],[-99.473167,45.571615],[-99.463394,45.57168],[-99.463391,45.57883],[-99.463368,45.586076],[-99.48177,45.585926],[-99.48384,45.585953],[-99.483885,45.57873],[-99.483911,45.577697]]]
+            },
+            "properties":{
+                "admin_level":"8",
+                "border_type":"city",
+                "boundary":"administrative",
+                "gnis:feature_id":"1267426",
+                "name":"Hosmer",
+                "place":"city",
+                "source":"TIGER/Line® 2008 Place Shapefiles (http://www.census.gov/geo/www/tiger/)",
+                "wikidata":"Q2442118",
+                "wikipedia":"en:Hosmer, South Dakota"
+            }
+        }"#,
+    )?;
+    index_writer.add_document(doc)?;
+    index_writer.commit()?;
+
+    let reader = index.reader()?;
+    let searcher = reader.searcher();
+    let field = schema.get_field("geometry").unwrap();
+    let query = SpatialQuery::from_bounds(field, [[-99.49, 45.56], [-99.45, 45.59]]);
+    let hits = searcher.search(&query, &TopDocs::with_limit(10).order_by_score())?;
+    for (_score, doc_address) in &hits {
+        let retrieved_doc: TantivyDocument = searcher.doc(*doc_address)?;
+        if let Some(field_value) = retrieved_doc.get_first(field) {
+            if let Some(geometry_box) = field_value.as_value().into_geometry() {
+                println!("Retrieved geometry: {:?}", geometry_box);
+            }
+        }
+    }
+    assert_eq!(hits.len(), 1);
+    Ok(())
+}

query-grammar/Cargo.toml

@@ -15,5 +15,5 @@ edition = "2024"
 nom = "7"
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.140"
-ordered-float = "5.0.0"
+ordered-float = { version = "5.0.0", features = ["serde"] }
 fnv = "1.0.7"

query-grammar/src/lib.rs

@@ -11,7 +11,7 @@ pub use crate::infallible::LenientError;
 pub use crate::occur::Occur;
 use crate::query_grammar::{parse_to_ast, parse_to_ast_lenient};
 pub use crate::user_input_ast::{
-    Delimiter, UserInputAst, UserInputBound, UserInputLeaf, UserInputLiteral,
+    Delimiter, SpatialPredicateKind, UserInputAst, UserInputBound, UserInputLeaf, UserInputLiteral,
 };
 
 #[derive(Debug, Serialize)]
@@ -58,6 +58,52 @@ mod tests {
         assert!(parse_query("title:").is_err());
     }
 
+    #[test]
+    fn test_spatial_intersects() {
+        // Test simple case first.
+        let ast = parse_query(r#"geo:$intersects(1.0 2.0, 3.0 4.0, 5.0 6.0)"#).unwrap();
+        let debug = format!("{ast:?}");
+        assert!(debug.contains("$intersects"), "got: {debug}");
+    }
+
+    #[test]
+    fn test_spatial_contains() {
+        let ast = parse_query(r#"geo:$contains(1.0 2.0, 3.0 4.0, 5.0 6.0)"#).unwrap();
+        let debug = format!("{ast:?}");
+        assert!(debug.contains("$contains"), "got: {debug}");
+    }
+
+    #[test]
+    fn test_spatial_join_debug() {
+        for (query, expected) in [
+            (
+                "geo:$within(1rad, $query(name:point))",
+                "\"geo\":$within(1rad, $query(name:point))",
+            ),
+            (
+                "geo:$between(1rad, 2rad, $query(a))",
+                "\"geo\":$between(1rad, 2rad, $query(a))",
+            ),
+            (
+                "geo:$intersects($query(name:x))",
+                "\"geo\":$intersects($query(name:x))",
+            ),
+            (
+                "geo:$contains($query(name:x))",
+                "\"geo\":$contains($query(name:x))",
+            ),
+            ("geo:$knn(3, 1 2)", "\"geo\":$knn(3, 1 2)"),
+        ] {
+            let debug = format!("{:?}", parse_query(query).unwrap());
+            assert!(debug.contains(expected), "{query} gives {debug}");
+        }
+    }
+
+    #[test]
+    fn test_spatial_too_few_coords() {
+        assert!(parse_query(r#"geo:$intersects(1.0 2.0, 3.0 4.0)"#).is_err());
+    }
+
     #[test]
     fn test_parse_query_lenient_wrong_query() {
         let (_, errors) = parse_query_lenient("title:");