From 81a44ac4e04528f379bb5181370599b8f85dd122 Mon Sep 17 00:00:00 2001
From: ThomasThelen <tommythelen@gmail.com>
Date: Mon, 29 Apr 2024 11:27:08 -0700
Subject: [PATCH] Add support for s2 indexing

---
 .../java/com/ontotext/trree/geosparql/GeoSparqlConfig.java   | 3 ++-
 .../ontotext/trree/geosparql/lucene/LuceneGeoIndexer.java    | 3 +++
 .../com/ontotext/trree/geosparql/util/GeoSparqlUtils.java    | 4 ++++
 .../ontotext/trree/geosparql/util/GeoSparqlUtilsTest.java    | 5 +++++
 4 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/main/java/com/ontotext/trree/geosparql/GeoSparqlConfig.java b/src/main/java/com/ontotext/trree/geosparql/GeoSparqlConfig.java
index a712d9d..9e1bc96 100644
--- a/src/main/java/com/ontotext/trree/geosparql/GeoSparqlConfig.java
+++ b/src/main/java/com/ontotext/trree/geosparql/GeoSparqlConfig.java
@@ -16,7 +16,8 @@
 public class GeoSparqlConfig {
     public enum PrefixTree {
         GEOHASH(SimpleValueFactory.getInstance().createLiteral("geohash")),
-        QUAD(SimpleValueFactory.getInstance().createLiteral("quad"));
+        QUAD(SimpleValueFactory.getInstance().createLiteral("quad")),
+        S2(SimpleValueFactory.getInstance().createLiteral("s2"));
 
         private Literal literal;
 
diff --git a/src/main/java/com/ontotext/trree/geosparql/lucene/LuceneGeoIndexer.java b/src/main/java/com/ontotext/trree/geosparql/lucene/LuceneGeoIndexer.java
index 75aed7d..45d4bbe 100644
--- a/src/main/java/com/ontotext/trree/geosparql/lucene/LuceneGeoIndexer.java
+++ b/src/main/java/com/ontotext/trree/geosparql/lucene/LuceneGeoIndexer.java
@@ -14,6 +14,7 @@
 import org.apache.lucene.spatial.prefix.tree.GeohashPrefixTree;
 import org.apache.lucene.spatial.prefix.tree.QuadPrefixTree;
 import org.apache.lucene.spatial.prefix.tree.SpatialPrefixTree;
+import org.apache.lucene.spatial.prefix.tree.S2PrefixTree;
 import org.apache.lucene.spatial.query.SpatialArgs;
 import org.apache.lucene.spatial.query.SpatialOperation;
 import org.apache.lucene.spatial.serialized.SerializedDVStrategy;
@@ -100,6 +101,8 @@ public void initSettings() {
 			grid = new QuadPrefixTree(ctx, precision);
 		} else if (prefixTree == GeoSparqlConfig.PrefixTree.GEOHASH) {
 			grid = new GeohashPrefixTree(ctx, precision);
+		} else if (prefixTree == GeoSparqlConfig.PrefixTree.S2) {
+			grid = new S2PrefixTree(ctx, precision);
 		} else {
 			throw new PluginException("Unexpected prefix tree type: " + prefixTree);
 		}
diff --git a/src/main/java/com/ontotext/trree/geosparql/util/GeoSparqlUtils.java b/src/main/java/com/ontotext/trree/geosparql/util/GeoSparqlUtils.java
index 95b369d..9d3ebf9 100644
--- a/src/main/java/com/ontotext/trree/geosparql/util/GeoSparqlUtils.java
+++ b/src/main/java/com/ontotext/trree/geosparql/util/GeoSparqlUtils.java
@@ -89,6 +89,10 @@ public static void validateParams(GeoSparqlConfig.PrefixTree prefixTree, int pre
                 if (precision <= 0 || precision > QuadPrefixTree.MAX_LEVELS_POSSIBLE) {
                     throw new PluginException(constructExceptionMessage(prefixTree));
                 }
+            case S2:
+                if (precision <= 0) {
+                    throw new PluginException(constructExceptionMessage(prefixTree));
+                }
         }
     }
 
diff --git a/src/test/java/com/ontotext/trree/geosparql/util/GeoSparqlUtilsTest.java b/src/test/java/com/ontotext/trree/geosparql/util/GeoSparqlUtilsTest.java
index 1f1d80b..73f8ac1 100644
--- a/src/test/java/com/ontotext/trree/geosparql/util/GeoSparqlUtilsTest.java
+++ b/src/test/java/com/ontotext/trree/geosparql/util/GeoSparqlUtilsTest.java
@@ -25,6 +25,11 @@ public void testCanSaveAndRead2() {
         test(false, GeoSparqlConfig.PrefixTree.GEOHASH, 7, GeoSparqlConfig.PrefixTree.QUAD, 12, false);
     }
 
+    @Test
+    public void testCanSaveAndRead3() {
+        test(false, GeoSparqlConfig.PrefixTree.S2, 7, GeoSparqlConfig.PrefixTree.S2, 12, false);
+    }
+
     private void test(boolean enabled, GeoSparqlConfig.PrefixTree prefixTree, int precision,
                       GeoSparqlConfig.PrefixTree currentPrefixTree, int currentPrecision, boolean ignoreErrors) {
         GeoSparqlConfig config1 = new GeoSparqlConfig();