zhangfengcdt commented on code in PR #2073:
URL: https://github.com/apache/sedona/pull/2073#discussion_r2193653763
##########
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:
Sure, I have added a test (test_toDf_preserves_columns_with_proper_types)
that specifically tests the toDf method to preserve the type and values.
--
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]
