paleolimbot commented on code in PR #768: URL: https://github.com/apache/sedona-db/pull/768#discussion_r3093678219
########## python/sedonadb/tests/test_sjoin_gpu.py: ########## @@ -0,0 +1,339 @@ +# 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. + +import json +import re +import warnings + +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +from sedonadb.testing import PostGIS, SedonaDB, random_geometry +from shapely.geometry import Point + + +def _enable_gpu_or_skip(eng_sedonadb): + try: + eng_sedonadb.con.sql("SET gpu.enable = true").execute() + except Exception as e: + pytest.skip(f"GPU runtime is not available: {e}") + + [email protected]( + "join_type", ["INNER JOIN", "LEFT OUTER JOIN", "RIGHT OUTER JOIN"] +) [email protected]( + "on", + [ + "ST_Intersects(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Within(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Contains(sjoin_polygon.geometry, sjoin_point.geometry)", + ], +) +def test_spatial_join_gpu(join_type, on): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) + df_point = random_geometry("Point", 100, seed=42) + df_polygon = random_geometry( + "Polygon", 100, hole_rate=0.5, num_vertices=(2, 10), seed=43 + ) + + eng_sedonadb.create_table_arrow("sjoin_point", df_point) + eng_sedonadb.create_table_arrow("sjoin_polygon", df_polygon) + eng_postgis.create_table_arrow("sjoin_point", df_point) + eng_postgis.create_table_arrow("sjoin_polygon", df_polygon) + + sql = f""" + SELECT sjoin_point.id id0, sjoin_polygon.id id1 + FROM sjoin_point {join_type} sjoin_polygon + ON {on} + ORDER BY id0, id1 + """ + + sedonadb_results = eng_sedonadb.execute_and_collect(sql).to_pandas() + assert len(sedonadb_results) > 0 + eng_postgis.assert_query_result(sql, sedonadb_results) + + +def _plan_text(df): + query_plan = df.to_pandas() + return "\n".join(query_plan.iloc[:, 1].astype(str).tolist()) + + +def _spatial_join_side_file_names(plan_text): + """Extract the left/right parquet file names used by `SpatialJoinExec`. + + Example input: + SpatialJoinExec: join_type=Inner, on=ST_intersects(geo_right@0, geo_left@0) + ProjectionExec: expr=[geometry@0 as geo_right] + DataSourceExec: file_groups={1 group: [[.../natural-earth_countries_geo.parquet]]}, projection=[geometry], file_type=parquet + ProbeShuffleExec: partitioning=RoundRobinBatch(1) + ProjectionExec: expr=[geometry@0 as geo_left] + DataSourceExec: file_groups={1 group: [[.../natural-earth_cities_geo.parquet]]}, projection=[geometry], file_type=parquet + + Example output: + ["natural-earth_countries_geo", "natural-earth_cities_geo"] + """ + spatial_join_idx = plan_text.find("SpatialJoinExec:") + assert spatial_join_idx != -1, plan_text + + file_names = re.findall( + r"DataSourceExec:.*?/([^/\]]+)\.parquet", plan_text[spatial_join_idx:] + ) + assert len(file_names) >= 2, plan_text + return file_names[:2] + + [email protected]( + "join_type", + [ + "LEFT SEMI JOIN", + "LEFT ANTI JOIN", + "RIGHT SEMI JOIN", + "RIGHT ANTI JOIN", + ], +) [email protected]( + "on", + [ + "ST_Intersects(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Within(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Contains(sjoin_polygon.geometry, sjoin_point.geometry)", + ], +) +def test_spatial_join_semi_anti_gpu(join_type, on): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) + options = json.dumps( + { + "geom_type": "Point", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 42, + } + ) + df_point = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + options = json.dumps( + { + "geom_type": "Polygon", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 43, + } + ) + df_polygon = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + eng_sedonadb.create_table_arrow("sjoin_point", df_point) + eng_sedonadb.create_table_arrow("sjoin_polygon", df_polygon) + eng_postgis.create_table_arrow("sjoin_point", df_point) + eng_postgis.create_table_arrow("sjoin_polygon", df_polygon) + + is_left = join_type.startswith("LEFT") + is_semi = "SEMI" in join_type + + if is_left: + sedona_sql = f""" + SELECT sjoin_point.id id0 + FROM sjoin_point {join_type} sjoin_polygon + ON {on} + ORDER BY id0 + """ + exists = f"EXISTS (SELECT 1 FROM sjoin_polygon WHERE {on})" + where = exists if is_semi else f"NOT {exists}" + postgis_sql = f""" + SELECT sjoin_point.id id0 + FROM sjoin_point + WHERE {where} + ORDER BY id0 + """ + else: + sedona_sql = f""" + SELECT sjoin_polygon.id id1 + FROM sjoin_point {join_type} sjoin_polygon + ON {on} + ORDER BY id1 + """ + exists = f"EXISTS (SELECT 1 FROM sjoin_point WHERE {on})" + where = exists if is_semi else f"NOT {exists}" + postgis_sql = f""" + SELECT sjoin_polygon.id id1 + FROM sjoin_polygon + WHERE {where} + ORDER BY id1 + """ + + sedonadb_results = eng_sedonadb.execute_and_collect(sedona_sql).to_pandas() + assert len(sedonadb_results) > 0 + eng_postgis.assert_query_result(postgis_sql, sedonadb_results) + + [email protected]( + "outer", + ["point", "polygon"], +) [email protected]( + "on", + [ + "ST_Intersects(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Within(sjoin_point.geometry, sjoin_polygon.geometry)", + ], +) +def test_spatial_mark_join_via_correlated_exists_gpu(outer, on): Review Comment: I don't think you need this one (it's testing query planning) ########## python/sedonadb/tests/test_sjoin_gpu.py: ########## @@ -0,0 +1,339 @@ +# 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. + +import json +import re +import warnings + +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +from sedonadb.testing import PostGIS, SedonaDB, random_geometry +from shapely.geometry import Point + + +def _enable_gpu_or_skip(eng_sedonadb): + try: + eng_sedonadb.con.sql("SET gpu.enable = true").execute() + except Exception as e: + pytest.skip(f"GPU runtime is not available: {e}") + + [email protected]( + "join_type", ["INNER JOIN", "LEFT OUTER JOIN", "RIGHT OUTER JOIN"] +) [email protected]( + "on", + [ + "ST_Intersects(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Within(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Contains(sjoin_polygon.geometry, sjoin_point.geometry)", + ], +) +def test_spatial_join_gpu(join_type, on): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) + df_point = random_geometry("Point", 100, seed=42) + df_polygon = random_geometry( + "Polygon", 100, hole_rate=0.5, num_vertices=(2, 10), seed=43 + ) + + eng_sedonadb.create_table_arrow("sjoin_point", df_point) + eng_sedonadb.create_table_arrow("sjoin_polygon", df_polygon) + eng_postgis.create_table_arrow("sjoin_point", df_point) + eng_postgis.create_table_arrow("sjoin_polygon", df_polygon) + + sql = f""" + SELECT sjoin_point.id id0, sjoin_polygon.id id1 + FROM sjoin_point {join_type} sjoin_polygon + ON {on} + ORDER BY id0, id1 + """ + + sedonadb_results = eng_sedonadb.execute_and_collect(sql).to_pandas() + assert len(sedonadb_results) > 0 + eng_postgis.assert_query_result(sql, sedonadb_results) + + +def _plan_text(df): + query_plan = df.to_pandas() + return "\n".join(query_plan.iloc[:, 1].astype(str).tolist()) + + +def _spatial_join_side_file_names(plan_text): + """Extract the left/right parquet file names used by `SpatialJoinExec`. + + Example input: + SpatialJoinExec: join_type=Inner, on=ST_intersects(geo_right@0, geo_left@0) + ProjectionExec: expr=[geometry@0 as geo_right] + DataSourceExec: file_groups={1 group: [[.../natural-earth_countries_geo.parquet]]}, projection=[geometry], file_type=parquet + ProbeShuffleExec: partitioning=RoundRobinBatch(1) + ProjectionExec: expr=[geometry@0 as geo_left] + DataSourceExec: file_groups={1 group: [[.../natural-earth_cities_geo.parquet]]}, projection=[geometry], file_type=parquet + + Example output: + ["natural-earth_countries_geo", "natural-earth_cities_geo"] + """ + spatial_join_idx = plan_text.find("SpatialJoinExec:") + assert spatial_join_idx != -1, plan_text + + file_names = re.findall( + r"DataSourceExec:.*?/([^/\]]+)\.parquet", plan_text[spatial_join_idx:] + ) + assert len(file_names) >= 2, plan_text + return file_names[:2] + + [email protected]( + "join_type", + [ + "LEFT SEMI JOIN", + "LEFT ANTI JOIN", + "RIGHT SEMI JOIN", + "RIGHT ANTI JOIN", + ], +) [email protected]( + "on", + [ + "ST_Intersects(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Within(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Contains(sjoin_polygon.geometry, sjoin_point.geometry)", + ], +) +def test_spatial_join_semi_anti_gpu(join_type, on): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) + options = json.dumps( + { + "geom_type": "Point", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 42, + } + ) + df_point = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + options = json.dumps( + { + "geom_type": "Polygon", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 43, + } + ) + df_polygon = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + eng_sedonadb.create_table_arrow("sjoin_point", df_point) + eng_sedonadb.create_table_arrow("sjoin_polygon", df_polygon) + eng_postgis.create_table_arrow("sjoin_point", df_point) + eng_postgis.create_table_arrow("sjoin_polygon", df_polygon) + + is_left = join_type.startswith("LEFT") + is_semi = "SEMI" in join_type + + if is_left: + sedona_sql = f""" + SELECT sjoin_point.id id0 + FROM sjoin_point {join_type} sjoin_polygon + ON {on} + ORDER BY id0 + """ + exists = f"EXISTS (SELECT 1 FROM sjoin_polygon WHERE {on})" + where = exists if is_semi else f"NOT {exists}" + postgis_sql = f""" + SELECT sjoin_point.id id0 + FROM sjoin_point + WHERE {where} + ORDER BY id0 + """ + else: + sedona_sql = f""" + SELECT sjoin_polygon.id id1 + FROM sjoin_point {join_type} sjoin_polygon + ON {on} + ORDER BY id1 + """ + exists = f"EXISTS (SELECT 1 FROM sjoin_point WHERE {on})" + where = exists if is_semi else f"NOT {exists}" + postgis_sql = f""" + SELECT sjoin_polygon.id id1 + FROM sjoin_polygon + WHERE {where} + ORDER BY id1 + """ + + sedonadb_results = eng_sedonadb.execute_and_collect(sedona_sql).to_pandas() + assert len(sedonadb_results) > 0 + eng_postgis.assert_query_result(postgis_sql, sedonadb_results) + + [email protected]( + "outer", + ["point", "polygon"], +) [email protected]( + "on", + [ + "ST_Intersects(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Within(sjoin_point.geometry, sjoin_polygon.geometry)", + ], +) +def test_spatial_mark_join_via_correlated_exists_gpu(outer, on): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) + options = json.dumps( + { + "geom_type": "Point", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 42, + } + ) + df_point = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + options = json.dumps( + { + "geom_type": "Polygon", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 43, + } + ) + df_polygon = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + eng_sedonadb.create_table_arrow("sjoin_point", df_point) + eng_sedonadb.create_table_arrow("sjoin_polygon", df_polygon) + eng_postgis.create_table_arrow("sjoin_point", df_point) + eng_postgis.create_table_arrow("sjoin_polygon", df_polygon) + + if outer == "point": + sql = f""" + SELECT sjoin_point.id id0 + FROM sjoin_point + WHERE sjoin_point.id = 1 OR EXISTS (SELECT 1 FROM sjoin_polygon WHERE {on}) + ORDER BY id0 + """ + else: + sql = f""" + SELECT sjoin_polygon.id id1, ST_AsBinary(sjoin_polygon.geometry) geom + FROM sjoin_polygon + WHERE sjoin_polygon.id = 1 OR EXISTS (SELECT 1 FROM sjoin_point WHERE {on}) + ORDER BY id1 + """ + + query_plan = eng_sedonadb.execute_and_collect(f"EXPLAIN {sql}").to_pandas() + plan_text = "\n".join(query_plan.iloc[:, 1].astype(str).tolist()) + assert any( + "SpatialJoinExec" in line and ("LeftMark" in line or "RightMark" in line) + for line in plan_text.splitlines() + ), plan_text + + sedonadb_results = eng_sedonadb.execute_and_collect(sql).to_pandas() + assert len(sedonadb_results) > 0 + eng_postgis.assert_query_result(sql, sedonadb_results) + + +def test_query_window_in_subquery_gpu(): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) Review Comment: I also don't think you need this one ########## python/sedonadb/tests/test_sjoin_gpu.py: ########## @@ -0,0 +1,339 @@ +# 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. + +import json +import re +import warnings + +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +from sedonadb.testing import PostGIS, SedonaDB, random_geometry +from shapely.geometry import Point + + +def _enable_gpu_or_skip(eng_sedonadb): + try: + eng_sedonadb.con.sql("SET gpu.enable = true").execute() + except Exception as e: + pytest.skip(f"GPU runtime is not available: {e}") + Review Comment: This works, although we have a `__features__` that I've been using for geography to skip tests in this way: https://github.com/apache/sedona-db/blob/742e851bd000d798fe3c09918a30a7cf044c7717/python/sedonadb/src/lib.rs#L56-L66 https://github.com/apache/sedona-db/blob/742e851bd000d798fe3c09918a30a7cf044c7717/python/sedonadb/tests/test_sjoin.py#L340-L341 ########## python/sedonadb/tests/test_sjoin_gpu.py: ########## @@ -0,0 +1,339 @@ +# 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. + +import json +import re +import warnings + +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +from sedonadb.testing import PostGIS, SedonaDB, random_geometry +from shapely.geometry import Point + + +def _enable_gpu_or_skip(eng_sedonadb): + try: + eng_sedonadb.con.sql("SET gpu.enable = true").execute() + except Exception as e: + pytest.skip(f"GPU runtime is not available: {e}") + + [email protected]( + "join_type", ["INNER JOIN", "LEFT OUTER JOIN", "RIGHT OUTER JOIN"] +) [email protected]( + "on", + [ + "ST_Intersects(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Within(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Contains(sjoin_polygon.geometry, sjoin_point.geometry)", + ], +) +def test_spatial_join_gpu(join_type, on): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) + df_point = random_geometry("Point", 100, seed=42) + df_polygon = random_geometry( + "Polygon", 100, hole_rate=0.5, num_vertices=(2, 10), seed=43 + ) + + eng_sedonadb.create_table_arrow("sjoin_point", df_point) + eng_sedonadb.create_table_arrow("sjoin_polygon", df_polygon) + eng_postgis.create_table_arrow("sjoin_point", df_point) + eng_postgis.create_table_arrow("sjoin_polygon", df_polygon) + + sql = f""" + SELECT sjoin_point.id id0, sjoin_polygon.id id1 + FROM sjoin_point {join_type} sjoin_polygon + ON {on} + ORDER BY id0, id1 + """ + + sedonadb_results = eng_sedonadb.execute_and_collect(sql).to_pandas() + assert len(sedonadb_results) > 0 + eng_postgis.assert_query_result(sql, sedonadb_results) + + +def _plan_text(df): + query_plan = df.to_pandas() + return "\n".join(query_plan.iloc[:, 1].astype(str).tolist()) + + +def _spatial_join_side_file_names(plan_text): + """Extract the left/right parquet file names used by `SpatialJoinExec`. + + Example input: + SpatialJoinExec: join_type=Inner, on=ST_intersects(geo_right@0, geo_left@0) + ProjectionExec: expr=[geometry@0 as geo_right] + DataSourceExec: file_groups={1 group: [[.../natural-earth_countries_geo.parquet]]}, projection=[geometry], file_type=parquet + ProbeShuffleExec: partitioning=RoundRobinBatch(1) + ProjectionExec: expr=[geometry@0 as geo_left] + DataSourceExec: file_groups={1 group: [[.../natural-earth_cities_geo.parquet]]}, projection=[geometry], file_type=parquet + + Example output: + ["natural-earth_countries_geo", "natural-earth_cities_geo"] + """ + spatial_join_idx = plan_text.find("SpatialJoinExec:") + assert spatial_join_idx != -1, plan_text + + file_names = re.findall( + r"DataSourceExec:.*?/([^/\]]+)\.parquet", plan_text[spatial_join_idx:] + ) + assert len(file_names) >= 2, plan_text + return file_names[:2] + + [email protected]( + "join_type", + [ + "LEFT SEMI JOIN", + "LEFT ANTI JOIN", + "RIGHT SEMI JOIN", + "RIGHT ANTI JOIN", + ], +) [email protected]( + "on", + [ + "ST_Intersects(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Within(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Contains(sjoin_polygon.geometry, sjoin_point.geometry)", + ], +) +def test_spatial_join_semi_anti_gpu(join_type, on): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) + options = json.dumps( + { + "geom_type": "Point", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 42, + } + ) + df_point = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + options = json.dumps( + { + "geom_type": "Polygon", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 43, + } + ) + df_polygon = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + eng_sedonadb.create_table_arrow("sjoin_point", df_point) + eng_sedonadb.create_table_arrow("sjoin_polygon", df_polygon) + eng_postgis.create_table_arrow("sjoin_point", df_point) + eng_postgis.create_table_arrow("sjoin_polygon", df_polygon) + + is_left = join_type.startswith("LEFT") + is_semi = "SEMI" in join_type + + if is_left: + sedona_sql = f""" + SELECT sjoin_point.id id0 + FROM sjoin_point {join_type} sjoin_polygon + ON {on} + ORDER BY id0 + """ + exists = f"EXISTS (SELECT 1 FROM sjoin_polygon WHERE {on})" + where = exists if is_semi else f"NOT {exists}" + postgis_sql = f""" + SELECT sjoin_point.id id0 + FROM sjoin_point + WHERE {where} + ORDER BY id0 + """ + else: + sedona_sql = f""" + SELECT sjoin_polygon.id id1 + FROM sjoin_point {join_type} sjoin_polygon + ON {on} + ORDER BY id1 + """ + exists = f"EXISTS (SELECT 1 FROM sjoin_point WHERE {on})" + where = exists if is_semi else f"NOT {exists}" + postgis_sql = f""" + SELECT sjoin_polygon.id id1 + FROM sjoin_polygon + WHERE {where} + ORDER BY id1 + """ + + sedonadb_results = eng_sedonadb.execute_and_collect(sedona_sql).to_pandas() + assert len(sedonadb_results) > 0 + eng_postgis.assert_query_result(postgis_sql, sedonadb_results) + + [email protected]( + "outer", + ["point", "polygon"], +) [email protected]( + "on", + [ + "ST_Intersects(sjoin_point.geometry, sjoin_polygon.geometry)", + "ST_Within(sjoin_point.geometry, sjoin_polygon.geometry)", + ], +) +def test_spatial_mark_join_via_correlated_exists_gpu(outer, on): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) + options = json.dumps( + { + "geom_type": "Point", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 42, + } + ) + df_point = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + options = json.dumps( + { + "geom_type": "Polygon", + "polygon_hole_rate": 0.5, + "num_parts_range": [2, 10], + "vertices_per_linestring_range": [2, 10], + "seed": 43, + } + ) + df_polygon = eng_sedonadb.execute_and_collect( + f"SELECT * FROM sd_random_geometry('{options}') LIMIT 100" + ) + eng_sedonadb.create_table_arrow("sjoin_point", df_point) + eng_sedonadb.create_table_arrow("sjoin_polygon", df_polygon) + eng_postgis.create_table_arrow("sjoin_point", df_point) + eng_postgis.create_table_arrow("sjoin_polygon", df_polygon) + + if outer == "point": + sql = f""" + SELECT sjoin_point.id id0 + FROM sjoin_point + WHERE sjoin_point.id = 1 OR EXISTS (SELECT 1 FROM sjoin_polygon WHERE {on}) + ORDER BY id0 + """ + else: + sql = f""" + SELECT sjoin_polygon.id id1, ST_AsBinary(sjoin_polygon.geometry) geom + FROM sjoin_polygon + WHERE sjoin_polygon.id = 1 OR EXISTS (SELECT 1 FROM sjoin_point WHERE {on}) + ORDER BY id1 + """ + + query_plan = eng_sedonadb.execute_and_collect(f"EXPLAIN {sql}").to_pandas() + plan_text = "\n".join(query_plan.iloc[:, 1].astype(str).tolist()) + assert any( + "SpatialJoinExec" in line and ("LeftMark" in line or "RightMark" in line) + for line in plan_text.splitlines() + ), plan_text + + sedonadb_results = eng_sedonadb.execute_and_collect(sql).to_pandas() + assert len(sedonadb_results) > 0 + eng_postgis.assert_query_result(sql, sedonadb_results) + + +def test_query_window_in_subquery_gpu(): + with ( + SedonaDB.create_or_skip() as eng_sedonadb, + PostGIS.create_or_skip() as eng_postgis, + ): + _enable_gpu_or_skip(eng_sedonadb) + + df_point = random_geometry("Point", 100, seed=100) + df_polygon = random_geometry( + "Polygon", 100, hole_rate=0.5, num_vertices=(2, 10), size=(50, 60), seed=999 + ) + + eng_sedonadb.create_table_arrow("sjoin_point", df_point) + eng_sedonadb.create_table_arrow("sjoin_polygon", df_polygon) + eng_postgis.create_table_arrow("sjoin_point", df_point) + eng_postgis.create_table_arrow("sjoin_polygon", df_polygon) + + sql = """ + SELECT id + FROM sjoin_point AS L + WHERE ST_Intersects(L.geometry, (SELECT R.geometry FROM sjoin_polygon AS R WHERE R.id = 1)) + ORDER BY id \ + """ + + query_plan = eng_sedonadb.execute_and_collect(f"EXPLAIN {sql}").to_pandas() + assert "SpatialJoinExec" in query_plan.iloc[1, 1] + + sedonadb_results = eng_sedonadb.execute_and_collect(sql).to_pandas() + assert len(sedonadb_results) > 0 + eng_postgis.assert_query_result(sql, sedonadb_results) + + +def test_spatial_join_with_pandas_metadata_gpu(con): + try: + con.sql("SET gpu.enable = true").execute() + except Exception as e: + pytest.skip(f"GPU runtime is not available: {e}") Review Comment: I don't think you need this one (more specific to the Python internals than anything about join execution) -- 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]
