Merge pull request #1 from djsdjs/patch-1

Update PolygonGeohashUDF.java
leovan · May 29, 2024 · 623062d · 623062d
2 parents 4cb846f + 918fb04
commit 623062d
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/src/main/java/tech/leovan/hive/udf/geo/PolygonAreaUDF.java b/src/main/java/tech/leovan/hive/udf/geo/PolygonAreaUDF.java
@@ -39,14 +39,14 @@ public Double evaluate(String polygonString, String polygonStringFormat) {
             return null;
         }
 
-        Geometry polygon = CTX.getShapeFactory().getGeometryFrom(shape);
-        if (!(polygon instanceof Polygon || polygon instanceof MultiPolygon)) {
+        Geometry geometry = CTX.getShapeFactory().getGeometryFrom(shape);
+        if (!(geometry instanceof Polygon || geometry instanceof MultiPolygon)) {
             return null;
         }
 
         try {
             MathTransform transform = CRS.findMathTransform(SOURCE_CRS, TARGET_CRS, false);
-            Geometry geometryMercator = JTS.transform(polygon, transform);
+            Geometry geometryMercator = JTS.transform(geometry, transform);
 
             return geometryMercator.getArea();
         } catch (Exception e) {

diff --git a/src/main/java/tech/leovan/hive/udf/geo/PolygonGeohashUDF.java b/src/main/java/tech/leovan/hive/udf/geo/PolygonGeohashUDF.java
@@ -42,11 +42,26 @@ public List<String> evaluate(
             return null;
         }
 
-        Geometry polygon = CTX.getShapeFactory().getGeometryFrom(shape);
-        if (!(polygon instanceof Polygon || polygon instanceof MultiPolygon)) {
+        Geometry geometry = CTX.getShapeFactory().getGeometryFrom(shape);
+        if (!(geometry instanceof Polygon || geometry instanceof MultiPolygon)) {
             return null;
         }
 
+        Set<String> geohashes = new HashSet<>();
+
+        if (geometry instanceof Polygon) {
+            geohashes = getGeoHashes((Polygon) geometry, precision);
+        } else {
+            for (int i = 0; i < geometry.getNumGeometries(); i++){
+                Geometry polygon = geometry.getGeometryN(i);
+                geohashes.addAll(getGeoHashes((Polygon) polygon, precision));
+            }
+        }
+
+        return new ArrayList<>(geohashes);
+    }
+
+    private Set<String> getGeoHashes(Polygon polygon, Integer precision){
         Point centroid = polygon.getCentroid();
 
         Set<String> geohashes = new HashSet<>();
@@ -55,8 +70,7 @@ public List<String> evaluate(
 
         while (!testingGeohashes.isEmpty()) {
             String geohash = testingGeohashes.poll();
-            Geometry geohashGeometry = CTX.getShapeFactory().getGeometryFrom(
-                    GeohashUtils.decodeBoundary(geohash, CTX));
+            Geometry geohashGeometry = CTX.getShapeFactory().getGeometryFrom(GeohashUtils.decodeBoundary(geohash, CTX));
 
             if (polygon.contains(geohashGeometry) || polygon.intersects(geohashGeometry)) {
                 geohashes.add(geohash);
@@ -72,7 +86,7 @@ public List<String> evaluate(
             }
         }
 
-        return new ArrayList<>(geohashes);
+        return geohashes;
     }
 
     public List<String> evaluate(