add Java test

apache · Jan 13, 2025 · 2c3642c · 2c3642c
1 parent 9bf2f61
commit 2c3642c
Showing 1 changed file with 23 additions and 8 deletions.
diff --git a/spark/common/src/test/java/org/apache/sedona/core/spatialRDD/PointRDDTest.java b/spark/common/src/test/java/org/apache/sedona/core/spatialRDD/PointRDDTest.java
@@ -21,7 +21,10 @@
 import static org.junit.Assert.assertEquals;
 
 import java.util.List;
+import java.util.Map;
+
 import org.apache.sedona.core.enums.IndexType;
+import org.apache.spark.api.java.JavaPairRDD;
 import org.junit.AfterClass;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -55,26 +58,22 @@ public void testConstructor() {
     spatialRDD.analyze();
     assertEquals(inputCount, spatialRDD.approximateTotalCount);
     assertEquals(inputBoundary, spatialRDD.boundaryEnvelope);
-    assert spatialRDD
-        .rawSpatialRDD
+    assert spatialRDD.rawSpatialRDD
         .take(9)
         .get(0)
         .getUserData()
         .equals("testattribute0\ttestattribute1\ttestattribute2");
-    assert spatialRDD
-        .rawSpatialRDD
+    assert spatialRDD.rawSpatialRDD
         .take(9)
         .get(2)
         .getUserData()
         .equals("testattribute0\ttestattribute1\ttestattribute2");
-    assert spatialRDD
-        .rawSpatialRDD
+    assert spatialRDD.rawSpatialRDD
         .take(9)
         .get(4)
         .getUserData()
         .equals("testattribute0\ttestattribute1\ttestattribute2");
-    assert spatialRDD
-        .rawSpatialRDD
+    assert spatialRDD.rawSpatialRDD
         .take(9)
         .get(8)
         .getUserData()
@@ -138,4 +137,20 @@ public void testBuildQuadtreeIndex() throws Exception {
       List<Point> result = spatialRDD.indexedRDD.take(1).get(0).query(spatialRDD.boundaryEnvelope);
     }
   }
+
+  @Test
+  public void testPartitionWithIds() throws Exception {
+    PointRDD spatialRDD = new PointRDD(sc, InputLocation, offset, splitter, true, numPartitions);
+    spatialRDD.analyze();
+    JavaPairRDD<Integer, Point> partitioned = spatialRDD.spatialPartitioningWithIds(gridType, 5);
+
+    // Make sure the output is partitioned
+    assertEquals(spatialRDD.getPartitioner().numPartitions(), partitioned.getNumPartitions());
+
+    Map<Integer, Long> partition_counts = partitioned.countByKey();
+    assertEquals(partitioned.getNumPartitions(), partition_counts.size());
+
+    // Make sure the count is the same
+    assertEquals(spatialRDD.rawSpatialRDD.count(), partitioned.count());
+  }
 }