Flatheadmill geo

Cargo.toml

@@ -56,6 +56,7 @@ itertools = "0.14.0"
 measure_time = "0.9.0"
 arc-swap = "1.5.0"
 bon = "3.3.1"
+i_triangle = "0.38.0"
 
 columnar = { version = "0.6", path = "./columnar", package = "tantivy-columnar" }
 sstable = { version = "0.6", path = "./sstable", package = "tantivy-sstable", optional = true }
@@ -70,6 +71,7 @@ futures-util = { version = "0.3.28", optional = true }
 futures-channel = { version = "0.3.28", optional = true }
 fnv = "1.0.7"
 typetag = "0.2.21"
+geo-types = "0.7.17"
 
 [target.'cfg(windows)'.dependencies]
 winapi = "0.3.9"

examples/geo_json.rs

@@ -0,0 +1,66 @@
+use geo_types::Point;
+use tantivy::collector::TopDocs;
+use tantivy::query::SpatialQuery;
+use tantivy::schema::{Schema, Value, SPATIAL, STORED, TEXT};
+use tantivy::spatial::point::GeoPoint;
+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::new(
+        field,
+        [
+            GeoPoint {
+                lon: -99.49,
+                lat: 45.56,
+            },
+            GeoPoint {
+                lon: -99.45,
+                lat: 45.59,
+            },
+        ],
+        tantivy::query::SpatialQueryType::Intersects,
+    );
+    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(())
+}

src/core/json_utils.rs

@@ -227,6 +227,9 @@ pub(crate) fn index_json_value<'a, V: Value<'a>>(
             ReferenceValueLeaf::IpAddr(_) => {
                 unimplemented!("IP address support in dynamic fields is not yet implemented")
             }
+            ReferenceValueLeaf::Geometry(_) => {
+                unimplemented!("Geometry support in dynamic fields is not implemented")
+            }
         },
         ReferenceValue::Array(elements) => {
             for val in elements {

src/fastfield/mod.rs

@@ -683,7 +683,7 @@ mod tests {
     }
 
     #[test]
-    fn test_datefastfield() -> crate::Result<()> {
+    fn test_datefastfield() {
         let mut schema_builder = Schema::builder();
         let date_field = schema_builder.add_date_field(
             "date",
@@ -697,22 +697,28 @@ mod tests {
         );
         let schema = schema_builder.build();
         let index = Index::create_in_ram(schema);
-        let mut index_writer = index.writer_for_tests()?;
+        let mut index_writer = index.writer_for_tests().unwrap();
         index_writer.set_merge_policy(Box::new(NoMergePolicy));
-        index_writer.add_document(doc!(
-            date_field => DateTime::from_u64(1i64.to_u64()),
-            multi_date_field => DateTime::from_u64(2i64.to_u64()),
-            multi_date_field => DateTime::from_u64(3i64.to_u64())
-        ))?;
-        index_writer.add_document(doc!(
-            date_field => DateTime::from_u64(4i64.to_u64())
-        ))?;
-        index_writer.add_document(doc!(
-            multi_date_field => DateTime::from_u64(5i64.to_u64()),
-            multi_date_field => DateTime::from_u64(6i64.to_u64())
-        ))?;
-        index_writer.commit()?;
-        let reader = index.reader()?;
+        index_writer
+            .add_document(doc!(
+                date_field => DateTime::from_u64(1i64.to_u64()),
+                multi_date_field => DateTime::from_u64(2i64.to_u64()),
+                multi_date_field => DateTime::from_u64(3i64.to_u64())
+            ))
+            .unwrap();
+        index_writer
+            .add_document(doc!(
+                date_field => DateTime::from_u64(4i64.to_u64())
+            ))
+            .unwrap();
+        index_writer
+            .add_document(doc!(
+                multi_date_field => DateTime::from_u64(5i64.to_u64()),
+                multi_date_field => DateTime::from_u64(6i64.to_u64())
+            ))
+            .unwrap();
+        index_writer.commit().unwrap();
+        let reader = index.reader().unwrap();
         let searcher = reader.searcher();
         assert_eq!(searcher.segment_readers().len(), 1);
         let segment_reader = searcher.segment_reader(0);
@@ -746,7 +752,6 @@ mod tests {
             assert_eq!(dates[0].into_timestamp_nanos(), 5i64);
             assert_eq!(dates[1].into_timestamp_nanos(), 6i64);
         }
-        Ok(())
     }
 
     #[test]

src/fastfield/writer.rs

@@ -189,6 +189,9 @@ impl FastFieldsWriter {
                             .record_str(doc_id, field_name, &token.text);
                     }
                 }
+                ReferenceValueLeaf::Geometry(_) => {
+                    panic!("Geometry fields should not be routed to fast field writer")
+                }
             },
             ReferenceValue::Array(val) => {
                 // TODO: Check this is the correct behaviour we want.
@@ -320,6 +323,9 @@ fn record_json_value_to_columnar_writer<'a, V: Value<'a>>(
                     "Pre-tokenized string support in dynamic fields is not yet implemented"
                 )
             }
+            ReferenceValueLeaf::Geometry(_) => {
+                unimplemented!("Geometry support in dynamic fields is not yet implemented")
+            }
         },
         ReferenceValue::Array(elements) => {
             for el in elements {

src/index/index_meta.rs

@@ -142,6 +142,7 @@ impl SegmentMeta {
             SegmentComponent::FastFields => ".fast".to_string(),
             SegmentComponent::FieldNorms => ".fieldnorm".to_string(),
             SegmentComponent::Delete => format!(".{}.del", self.delete_opstamp().unwrap_or(0)),
+            SegmentComponent::Spatial => ".spatial".to_string(),
         });
         PathBuf::from(path)
     }

src/index/segment_component.rs

@@ -28,12 +28,14 @@ pub enum SegmentComponent {
     /// Bitset describing which document of the segment is alive.
     /// (It was representing deleted docs but changed to represent alive docs from v0.17)
     Delete,
+    /// HUSH
+    Spatial,
 }
 
 impl SegmentComponent {
     /// Iterates through the components.
     pub fn iterator() -> slice::Iter<'static, SegmentComponent> {
-        static SEGMENT_COMPONENTS: [SegmentComponent; 8] = [
+        static SEGMENT_COMPONENTS: [SegmentComponent; 9] = [
             SegmentComponent::Postings,
             SegmentComponent::Positions,
             SegmentComponent::FastFields,
@@ -42,6 +44,7 @@ impl SegmentComponent {
             SegmentComponent::Store,
             SegmentComponent::TempStore,
             SegmentComponent::Delete,
+            SegmentComponent::Spatial,
         ];
         SEGMENT_COMPONENTS.iter()
     }

src/index/segment_reader.rs

@@ -14,6 +14,7 @@ use crate::index::{InvertedIndexReader, Segment, SegmentComponent, SegmentId};
 use crate::json_utils::json_path_sep_to_dot;
 use crate::schema::{Field, IndexRecordOption, Schema, Type};
 use crate::space_usage::SegmentSpaceUsage;
+use crate::spatial::reader::SpatialReaders;
 use crate::store::StoreReader;
 use crate::termdict::TermDictionary;
 use crate::{DocId, Opstamp};
@@ -43,6 +44,7 @@ pub struct SegmentReader {
     positions_composite: CompositeFile,
     fast_fields_readers: FastFieldReaders,
     fieldnorm_readers: FieldNormReaders,
+    spatial_readers: SpatialReaders,
 
     store_file: FileSlice,
     alive_bitset_opt: Option<AliveBitSet>,
@@ -92,6 +94,11 @@ impl SegmentReader {
         &self.fast_fields_readers
     }
 
+    /// HUSH
+    pub fn spatial_fields(&self) -> &SpatialReaders {
+        &self.spatial_readers
+    }
+
     /// Accessor to the `FacetReader` associated with a given `Field`.
     pub fn facet_reader(&self, field_name: &str) -> crate::Result<FacetReader> {
         let schema = self.schema();
@@ -173,6 +180,12 @@ impl SegmentReader {
         let fast_fields_readers = FastFieldReaders::open(fast_fields_data, schema.clone())?;
         let fieldnorm_data = segment.open_read(SegmentComponent::FieldNorms)?;
         let fieldnorm_readers = FieldNormReaders::open(fieldnorm_data)?;
+        let spatial_readers = if schema.contains_spatial_field() {
+            let spatial_data = segment.open_read(SegmentComponent::Spatial)?;
+            SpatialReaders::open(spatial_data)?
+        } else {
+            SpatialReaders::empty()
+        };
 
         let original_bitset = if segment.meta().has_deletes() {
             let alive_doc_file_slice = segment.open_read(SegmentComponent::Delete)?;
@@ -198,6 +211,7 @@ impl SegmentReader {
             postings_composite,
             fast_fields_readers,
             fieldnorm_readers,
+            spatial_readers,
             segment_id: segment.id(),
             delete_opstamp: segment.meta().delete_opstamp(),
             store_file,
@@ -460,6 +474,7 @@ impl SegmentReader {
             self.positions_composite.space_usage(),
             self.fast_fields_readers.space_usage(self.schema())?,
             self.fieldnorm_readers.space_usage(),
+            self.spatial_readers.space_usage(),
             self.get_store_reader(0)?.space_usage(),
             self.alive_bitset_opt
                 .as_ref()

src/indexer/merger.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+use std::io::{BufWriter, Write};
 use std::sync::Arc;
 
 use columnar::{
@@ -6,6 +8,7 @@ use columnar::{
 use common::ReadOnlyBitSet;
 use itertools::Itertools;
 use measure_time::debug_time;
+use tempfile::NamedTempFile;
 
 use crate::directory::WritePtr;
 use crate::docset::{DocSet, TERMINATED};
@@ -17,6 +20,8 @@ use crate::indexer::doc_id_mapping::{MappingType, SegmentDocIdMapping};
 use crate::indexer::SegmentSerializer;
 use crate::postings::{InvertedIndexSerializer, Postings, SegmentPostings};
 use crate::schema::{value_type_to_column_type, Field, FieldType, Schema};
+use crate::spatial::bkd::LeafPageIterator;
+use crate::spatial::triangle::Triangle;
 use crate::store::StoreWriter;
 use crate::termdict::{TermMerger, TermOrdinal};
 use crate::{DocAddress, DocId, InvertedIndexReader};
@@ -170,6 +175,7 @@ impl IndexMerger {
         let mut readers = vec![];
         for (segment, new_alive_bitset_opt) in segments.iter().zip(alive_bitset_opt) {
             if segment.meta().num_docs() > 0 {
+                dbg!("segment");
                 let reader =
                     SegmentReader::open_with_custom_alive_set(segment, new_alive_bitset_opt)?;
                 readers.push(reader);
@@ -520,6 +526,89 @@ impl IndexMerger {
         Ok(())
     }
 
+    fn write_spatial_fields(
+        &self,
+        serializer: &mut SegmentSerializer,
+        doc_id_mapping: &SegmentDocIdMapping,
+    ) -> crate::Result<()> {
+        /// We need to rebuild a BKD-tree based off the list of triangles.
+        ///
+        /// Because the data can be large, we do this by writing the sequence of triangles to
+        /// disk, and mmapping it as mutable slice, and calling the same code as what
+        /// is done for the segment serialization.
+        ///
+        /// The OS is in charge of deciding how to handle its page cache.
+        /// This is the same as what would have happened with swapping,
+        /// except by explicitly mapping the file, the OS is more likely to
+        /// swap, the memory will not be accounted as anonymous memory,
+        /// swap space is reserved etc.
+        use crate::spatial::bkd::Segment;
+
+        let Some(mut spatial_serializer) = serializer.extract_spatial_serializer() else {
+            // The schema does not contain any spatial field.
+            return Ok(());
+        };
+
+        let mut segment_mappings: Vec<Vec<Option<DocId>>> = Vec::new();
+        for reader in &self.readers {
+            let max_doc = reader.max_doc();
+            segment_mappings.push(vec![None; max_doc as usize]);
+        }
+        for (new_doc_id, old_doc_addr) in doc_id_mapping.iter_old_doc_addrs().enumerate() {
+            segment_mappings[old_doc_addr.segment_ord as usize][old_doc_addr.doc_id as usize] =
+                Some(new_doc_id as DocId);
+        }
+        let mut temp_files: HashMap<Field, NamedTempFile> = HashMap::new();
+
+        for (field, field_entry) in self.schema.fields() {
+            if matches!(field_entry.field_type(), FieldType::Spatial(_)) {
+                temp_files.insert(field, NamedTempFile::new()?);
+            }
+        }
+        for (segment_ord, reader) in self.readers.iter().enumerate() {
+            for (field, temp_file) in &mut temp_files {
+                let mut buf_temp_file = BufWriter::new(temp_file);
+                let spatial_readers = reader.spatial_fields();
+                let Some(spatial_reader) = spatial_readers.get_field(*field)? else {
+                    continue;
+                };
+                let segment = Segment::new(spatial_reader.get_bytes());
+                for triangle_result in LeafPageIterator::new(&segment) {
+                    let triangles = triangle_result?;
+                    for triangle in triangles {
+                        if let Some(new_doc_id) =
+                            segment_mappings[segment_ord][triangle.doc_id as usize]
+                        {
+                            // This is really just a temporary file, not meant to be portable, so we
+                            // use native endianness here.
+                            for &word in &triangle.words {
+                                buf_temp_file.write_all(&word.to_ne_bytes())?;
+                            }
+                            buf_temp_file.write_all(&new_doc_id.to_ne_bytes())?;
+                        }
+                    }
+                }
+                buf_temp_file.flush()?;
+                // No need to fsync here. This file is not here for persistency.
+            }
+        }
+        for (field, temp_file) in temp_files {
+            // Memory map the triangle file.
+            use memmap2::MmapOptions;
+            let mmap = unsafe { MmapOptions::new().map_mut(temp_file.as_file())? };
+            // Cast to &[Triangle] slice
+            let triangle_count = mmap.len() / std::mem::size_of::<Triangle>();
+            let triangles = unsafe {
+                std::slice::from_raw_parts_mut(mmap.as_ptr() as *mut Triangle, triangle_count)
+            };
+            // Get spatial writer and rebuild block kd-tree.
+            spatial_serializer.serialize_field(field, triangles)?;
+        }
+        spatial_serializer.close()?;
+
+        Ok(())
+    }
+
     /// Writes the merged segment by pushing information
     /// to the `SegmentSerializer`.
     ///
@@ -544,9 +633,10 @@ impl IndexMerger {
 
         debug!("write-storagefields");
         self.write_storable_fields(serializer.get_store_writer())?;
+        debug!("write-spatialfields");
+        self.write_spatial_fields(&mut serializer, &doc_id_mapping)?;
         debug!("write-fastfields");
         self.write_fast_fields(serializer.get_fast_field_write(), doc_id_mapping)?;
-
         debug!("close-serializer");
         serializer.close()?;
         Ok(self.max_doc)