| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| from shapely.geometry import Polygon |
| from tests import csv_point_input_location, union_polygon_input_location |
| from tests.test_base import TestBase |
| |
| |
| class TestConstructors(TestBase): |
| |
| def test_st_extent(self): |
| df = self.spark.sql(""" |
| SELECT ST_GeomFromWKT(wkt) AS geom |
| FROM VALUES ('POINT (1 2)'), ('POINT (4 5)'), ('LINESTRING (-3 0, 0 0)') AS t(wkt) |
| """) |
| df.createOrReplaceTempView("extent_input") |
| bbox = self.spark.sql("SELECT ST_Extent(geom) FROM extent_input").first()[0] |
| assert bbox.xmin == -3.0 |
| assert bbox.ymin == 0.0 |
| assert bbox.xmax == 4.0 |
| assert bbox.ymax == 5.0 |
| |
| def test_st_extent_returns_null_for_empty_input(self): |
| df = self.spark.sql(""" |
| SELECT ST_GeomFromWKT(wkt) AS geom |
| FROM VALUES (CAST(NULL AS STRING)), ('POINT EMPTY') AS t(wkt) |
| """) |
| df.createOrReplaceTempView("empty_extent_input") |
| result = self.spark.sql( |
| "SELECT ST_Extent(geom) FROM empty_extent_input" |
| ).first()[0] |
| assert result is None |
| |
| def test_st_envelope_aggr(self): |
| point_csv_df = ( |
| self.spark.read.format("csv") |
| .option("delimiter", ",") |
| .option("header", "false") |
| .load(csv_point_input_location) |
| ) |
| |
| point_csv_df.createOrReplaceTempView("pointtable") |
| point_df = self.spark.sql( |
| "select ST_Point(cast(pointtable._c0 as Decimal(24,20)), cast(pointtable._c1 as Decimal(24,20))) as arealandmark from pointtable" |
| ) |
| point_df.createOrReplaceTempView("pointdf") |
| boundary = self.spark.sql( |
| "select ST_Envelope_Aggr(pointdf.arealandmark) from pointdf" |
| ) |
| |
| coordinates = [ |
| (1.1, 101.1), |
| (1.1, 1100.1), |
| (1000.1, 1100.1), |
| (1000.1, 101.1), |
| (1.1, 101.1), |
| ] |
| |
| polygon = Polygon(coordinates) |
| |
| assert boundary.take(1)[0][0] == polygon |
| |
| def test_st_union_aggr(self): |
| polygon_csv_df = ( |
| self.spark.read.format("csv") |
| .option("delimiter", ",") |
| .option("header", "false") |
| .load(union_polygon_input_location) |
| ) |
| |
| polygon_csv_df.createOrReplaceTempView("polygontable") |
| polygon_csv_df.show() |
| polygon_df = self.spark.sql( |
| "select ST_PolygonFromEnvelope(cast(polygontable._c0 as Decimal(24,20)),cast(polygontable._c1 as Decimal(24,20)), cast(polygontable._c2 as Decimal(24,20)), cast(polygontable._c3 as Decimal(24,20))) as polygonshape from polygontable" |
| ) |
| polygon_df.createOrReplaceTempView("polygondf") |
| polygon_df.show() |
| union = self.spark.sql( |
| "select ST_Union_Aggr(polygondf.polygonshape) from polygondf" |
| ) |
| |
| assert union.take(1)[0][0].area == 10100 |
| |
| def test_st_collect_aggr_points(self): |
| self.spark.sql(""" |
| SELECT explode(array( |
| ST_GeomFromWKT('POINT(1 2)'), |
| ST_GeomFromWKT('POINT(3 4)'), |
| ST_GeomFromWKT('POINT(5 6)') |
| )) AS geom |
| """).createOrReplaceTempView("points_table") |
| |
| result = self.spark.sql("SELECT ST_Collect_Agg(geom) FROM points_table").take( |
| 1 |
| )[0][0] |
| |
| assert result.geom_type == "MultiPoint" |
| assert len(result.geoms) == 3 |
| |
| def test_st_collect_aggr_polygons(self): |
| self.spark.sql(""" |
| SELECT explode(array( |
| ST_GeomFromWKT('POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))'), |
| ST_GeomFromWKT('POLYGON((2 2, 3 2, 3 3, 2 3, 2 2))') |
| )) AS geom |
| """).createOrReplaceTempView("polygons_table") |
| |
| result = self.spark.sql("SELECT ST_Collect_Agg(geom) FROM polygons_table").take( |
| 1 |
| )[0][0] |
| |
| assert result.geom_type == "MultiPolygon" |
| assert len(result.geoms) == 2 |
| assert result.area == 2.0 |
| |
| def test_st_collect_aggr_mixed_types(self): |
| self.spark.sql(""" |
| SELECT explode(array( |
| ST_GeomFromWKT('POINT(1 2)'), |
| ST_GeomFromWKT('LINESTRING(0 0, 1 1)'), |
| ST_GeomFromWKT('POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))') |
| )) AS geom |
| """).createOrReplaceTempView("mixed_geom_table") |
| |
| result = self.spark.sql( |
| "SELECT ST_Collect_Agg(geom) FROM mixed_geom_table" |
| ).take(1)[0][0] |
| |
| assert result.geom_type == "GeometryCollection" |
| assert len(result.geoms) == 3 |
| |
| def test_st_collect_aggr_preserves_duplicates(self): |
| # Test that ST_Collect_Agg keeps duplicate geometries (unlike ST_Union_Aggr) |
| self.spark.sql(""" |
| SELECT explode(array( |
| ST_GeomFromWKT('POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))'), |
| ST_GeomFromWKT('POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))') |
| )) AS geom |
| """).createOrReplaceTempView("duplicate_polygons_table") |
| |
| result = self.spark.sql( |
| "SELECT ST_Collect_Agg(geom) FROM duplicate_polygons_table" |
| ).take(1)[0][0] |
| |
| # ST_Collect_Agg should preserve both polygons |
| assert len(result.geoms) == 2 |
| # Area should be 2 because it doesn't merge overlapping areas |
| assert result.area == 2.0 |
| |
| # Test aliases for *_Aggr functions with *_Agg suffix |
| def test_st_envelope_agg_alias(self): |
| self.spark.sql(""" |
| SELECT explode(array( |
| ST_GeomFromWKT('POINT(1.1 101.1)'), |
| ST_GeomFromWKT('POINT(1.1 1100.1)'), |
| ST_GeomFromWKT('POINT(1000.1 1100.1)'), |
| ST_GeomFromWKT('POINT(1000.1 101.1)') |
| )) AS arealandmark |
| """).createOrReplaceTempView("pointdf_alias") |
| |
| boundary = self.spark.sql( |
| "SELECT ST_Envelope_Agg(pointdf_alias.arealandmark) FROM pointdf_alias" |
| ) |
| |
| coordinates = [ |
| (1.1, 101.1), |
| (1.1, 1100.1), |
| (1000.1, 1100.1), |
| (1000.1, 101.1), |
| (1.1, 101.1), |
| ] |
| |
| polygon = Polygon(coordinates) |
| assert boundary.take(1)[0][0].equals(polygon) |
| |
| def test_st_intersection_agg_alias(self): |
| self.spark.sql(""" |
| SELECT explode(array( |
| ST_GeomFromWKT('POLYGON((0 0, 4 0, 4 4, 0 4, 0 0))'), |
| ST_GeomFromWKT('POLYGON((2 2, 6 2, 6 6, 2 6, 2 2))') |
| )) AS countyshape |
| """).createOrReplaceTempView("polygondf_alias") |
| |
| intersection = self.spark.sql( |
| "SELECT ST_Intersection_Agg(polygondf_alias.countyshape) FROM polygondf_alias" |
| ) |
| |
| result = intersection.take(1)[0][0] |
| # The intersection of the two polygons should be a square from (2,2) to (4,4) with area 4 |
| assert result.area == 4.0 |
| |
| def test_st_union_agg_alias(self): |
| self.spark.sql(""" |
| SELECT explode(array( |
| ST_GeomFromWKT('POLYGON((0 0, 2 0, 2 2, 0 2, 0 0))'), |
| ST_GeomFromWKT('POLYGON((1 1, 3 1, 3 3, 1 3, 1 1))') |
| )) AS countyshape |
| """).createOrReplaceTempView("polygondf_union_alias") |
| |
| union = self.spark.sql( |
| "SELECT ST_Union_Agg(polygondf_union_alias.countyshape) FROM polygondf_union_alias" |
| ) |
| |
| result = union.take(1)[0][0] |
| # Two overlapping 2x2 squares with 1x1 overlap: area = 4 + 4 - 1 = 7 |
| assert result.area == 7.0 |
