jiayuasu commented on code in PR #2073:
URL: https://github.com/apache/sedona/pull/2073#discussion_r2193631743
##########
python/tests/sql/test_structured_adapter.py:
##########
@@ -80,3 +85,124 @@ def test_spatial_partitioned_write(self):
out = td + "/out"
partitioned_df.write.format("geoparquet").save(out)
assert len(glob.glob(out + "/*.parquet")) == n_spatial_partitions
+
+ def test_build_index_and_range_query_with_polygons(self):
+ # Create a spatial DataFrame with polygons
+ polygons_data = [
+ (1, "POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))"),
+ (2, "POLYGON((1 1, 2 1, 2 2, 1 2, 1 1))"),
+ (3, "POLYGON((2 2, 3 2, 3 3, 2 3, 2 2))"),
+ (4, "POLYGON((3 3, 4 3, 4 4, 3 4, 3 3))"),
+ (5, "POLYGON((4 4, 5 4, 5 5, 4 5, 4 4))"),
+ ]
+
+ df = self.spark.createDataFrame(polygons_data, ["id", "wkt"])
+ spatial_df = df.withColumn("geometry", expr("ST_GeomFromWKT(wkt)"))
+
+ # Convert to SpatialRDD
+ spatial_rdd = StructuredAdapter.toSpatialRdd(spatial_df, "geometry")
+
+ # Build index on the spatial RDD
+ spatial_rdd.buildIndex(IndexType.RTREE, False)
+
+ query_point = Point(2.2, 2.2)
+
+ # Perform range query
+ query_result = RangeQuery.SpatialRangeQuery(
+ spatial_rdd, query_point, True, True
+ )
+
+ # Assertions
+ result_count = query_result.count()
+
+ assert result_count >= 0, f"Expected at least one result, got
{result_count}"
+
+
+def test_build_index_and_range_query_with_points(self):
+ # Create a spatial DataFrame with points
+ points_data = [
+ (1, "POINT(0 0)"),
+ (2, "POINT(1 1)"),
+ (3, "POINT(2 2)"),
+ (4, "POINT(3 3)"),
+ (5, "POINT(4 4)"),
+ ]
+
+ df = self.spark.createDataFrame(points_data, ["id", "wkt"])
+ spatial_df = df.withColumn("geometry", expr("ST_GeomFromWKT(wkt)"))
+
+ # Convert to SpatialRDD
+ spatial_rdd = StructuredAdapter.toSpatialRdd(spatial_df, "geometry")
+
+ # Build index on the spatial RDD
+ spatial_rdd.buildIndex(IndexType.RTREE, False)
+
+ query_window = Point(2.0, 2.0).buffer(1.0)
+
+ # Perform range query
+ query_result = RangeQuery.SpatialRangeQuery(spatial_rdd, query_window,
True, True)
+
+ # Assertions
+ result_count = query_result.count()
+ assert result_count > 0, f"Expected at least one result, got
{result_count}"
+
+
+def test_build_index_and_range_query_with_linestrings(self):
+ # Create a spatial DataFrame with linestrings
+ linestrings_data = [
+ (1, "LINESTRING(0 0, 1 1)"),
+ (2, "LINESTRING(1 1, 2 2)"),
+ (3, "LINESTRING(2 2, 3 3)"),
+ (4, "LINESTRING(3 3, 4 4)"),
+ (5, "LINESTRING(4 4, 5 5)"),
+ ]
+
+ df = self.spark.createDataFrame(linestrings_data, ["id", "wkt"])
+ spatial_df = df.withColumn("geometry", expr("ST_GeomFromWKT(wkt)"))
+
+ # Convert to SpatialRDD
+ spatial_rdd = StructuredAdapter.toSpatialRdd(spatial_df, "geometry")
+
+ # Build index on the spatial RDD
+ spatial_rdd.buildIndex(IndexType.RTREE, False)
+
+ query_window = Point(2.0, 2.0).buffer(0.5)
+
+ # Perform range query
+ query_result = RangeQuery.SpatialRangeQuery(spatial_rdd, query_window,
True, True)
+
+ # Assertions
+ result_count = query_result.count()
+ assert result_count > 0, f"Expected at least one result, got
{result_count}"
+
+
+def test_build_index_and_range_query_with_mixed_geometries(self):
+ # Create a spatial DataFrame with mixed geometry types
+ mixed_data = [
+ (1, "POINT(0 0)"),
+ (2, "LINESTRING(1 1, 2 2)"),
+ (3, "POLYGON((2 2, 3 2, 3 3, 2 3, 2 2))"),
+ (4, "MULTIPOINT((3 3), (3.1 3.1))"),
+ (
+ 5,
+ "MULTIPOLYGON(((4 4, 5 4, 5 5, 4 5, 4 4)), ((4.1 4.1, 4.2 4.1, 4.2
4.2, 4.1 4.2, 4.1 4.1)))",
+ ),
+ ]
+
+ df = self.spark.createDataFrame(mixed_data, ["id", "wkt"])
+ spatial_df = df.withColumn("geometry", expr("ST_GeomFromWKT(wkt)"))
+
+ # Convert to SpatialRDD
+ spatial_rdd = StructuredAdapter.toSpatialRdd(spatial_df, "geometry")
Review Comment:
Please add tests to make sure the StructuredAdapter.toDf can get the exact
same DF back including the non-spatial columns with correct values and types.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
