diff --git a/common/src/main/java/org/apache/sedona/common/Predicates.java b/common/src/main/java/org/apache/sedona/common/Predicates.java index 0511f4a2c49..92e6b656104 100644 --- a/common/src/main/java/org/apache/sedona/common/Predicates.java +++ b/common/src/main/java/org/apache/sedona/common/Predicates.java @@ -262,4 +262,23 @@ public static boolean knn( Geometry leftGeometry, Geometry rightGeometry, int k, boolean useSpheroid) { throw new UnsupportedOperationException("KNN predicate is not supported"); } + + public static boolean knn( + Geometry leftGeometry, + Geometry rightGeometry, + int k, + boolean useSpheroid, + boolean includeTies) { + throw new UnsupportedOperationException("KNN predicate is not supported"); + } + + public static boolean knn( + Geometry leftGeometry, + Geometry rightGeometry, + int k, + boolean useSpheroid, + boolean includeTies, + boolean exclusive) { + throw new UnsupportedOperationException("KNN predicate is not supported"); + } } diff --git a/docs/api/sql/NearestNeighbourSearching.md b/docs/api/sql/NearestNeighbourSearching.md index 26712a07a18..db95ca299fd 100644 --- a/docs/api/sql/NearestNeighbourSearching.md +++ b/docs/api/sql/NearestNeighbourSearching.md @@ -24,21 +24,43 @@ Sedona supports nearest-neighbour searching on geospatial data by providing a ge Introduction: join operation to find the k-nearest neighbors of a point or region in a spatial dataset. -Format: `ST_KNN(R: Table, S: Table, k: Integer, use_spheroid: Boolean)` +Formats: -Where R is the queries side table and S is the object side table, K is the number of neighbors. use_spheroid is a boolean value that determines whether to use the spheroid distance or not. +```sql +ST_KNN(query_geometry, object_geometry, k) +ST_KNN(query_geometry, object_geometry, k, use_spheroid) +ST_KNN(query_geometry, object_geometry, k, use_spheroid, include_ties) +ST_KNN(query_geometry, object_geometry, k, use_spheroid, include_ties, exclusive) +``` + +`query_geometry` belongs to the query-side table, `object_geometry` belongs to +the candidate-side table, and `k` is the number of neighbours. `use_spheroid` +selects spheroid distance instead of planar distance. + +The optional `include_ties` argument controls ties for this query. It overrides +the `spark.sedona.join.knn.includeTieBreakers` session setting when supplied. +The optional `exclusive` argument excludes candidates that are topologically +equal to the query geometry before nearest neighbours are ranked. For example, +the two line strings `LINESTRING (0 0, 1 0)` and +`LINESTRING (1 0, 0 0)` are equal and are excluded from one another by +`exclusive = true`. + +The three- and four-argument forms retain their existing behavior. To use +`exclusive`, both query-local Boolean arguments must be supplied. Queries side table contains geometries that are used to find the k-nearest neighbors in the object side table. When either queries or objects data contain non-point data (geometries), we take the centroid of each geometry. -In case there are ties in the distance, the result will include all the tied geometries only when the following sedona config is set to true: +With the three- or four-argument form, tied geometries are included only when +the following Sedona configuration is set to true: **Note for Inner Join:** - The `ST_KNN` join only supports left inner join. - It returns only pairs where there is at least one matching neighbor within the k nearest neighbors. - If a query point has no valid neighbor (e.g., because k is too large), it is excluded from the result. +- Null and empty geometries do not participate in a match. ``` spark.sedona.join.knn.includeTieBreakers=true diff --git a/docs/api/sql/NearestNeighbourSearching.zh.md b/docs/api/sql/NearestNeighbourSearching.zh.md index 15909cb58ef..ddcfc9b2fe5 100644 --- a/docs/api/sql/NearestNeighbourSearching.zh.md +++ b/docs/api/sql/NearestNeighbourSearching.zh.md @@ -24,21 +24,40 @@ Sedona 通过提供地理空间 k 近邻（kNN）连接方法来支持对地理 简介：用于在空间数据集中查找一个点或区域的 k 个最近邻的连接操作。 -格式：`ST_KNN(R: Table, S: Table, k: Integer, use_spheroid: Boolean)` +格式： -其中 R 是查询侧表，S 是对象侧表，K 是邻居数量。use_spheroid 是一个布尔值，决定是否使用椭球距离。 +```sql +ST_KNN(query_geometry, object_geometry, k) +ST_KNN(query_geometry, object_geometry, k, use_spheroid) +ST_KNN(query_geometry, object_geometry, k, use_spheroid, include_ties) +ST_KNN(query_geometry, object_geometry, k, use_spheroid, include_ties, exclusive) +``` + +`query_geometry` 属于查询侧表，`object_geometry` 属于候选对象侧表，`k` +是邻居数量。`use_spheroid` 用于选择椭球距离而不是平面距离。 + +可选参数 `include_ties` 仅控制当前查询的并列结果；提供该参数时，它会覆盖 +会话配置 `spark.sedona.join.knn.includeTieBreakers`。可选参数 `exclusive` +会在近邻排序之前排除与查询几何拓扑相等的候选对象。例如， +`LINESTRING (0 0, 1 0)` 和 `LINESTRING (1 0, 0 0)` 拓扑相等，因此在 +`exclusive = true` 时会互相排除。 + +三个和四个参数的形式保持原有行为。若要使用 `exclusive`，必须同时提供两个 +查询级布尔参数。 查询侧表包含用于在对象侧表中查找 k 近邻的几何对象。 当查询数据或对象数据中存在非点几何（其他几何类型）时，会取每个几何对象的质心。 -当距离上出现并列时，只有在下面的 sedona 配置被设置为 true 时，结果才会包含所有并列的几何对象： +使用三个或四个参数的形式时，只有在下面的 Sedona 配置被设置为 true 时， +结果才会包含所有并列的几何对象： **关于内连接的说明：** - `ST_KNN` 连接仅支持 left inner join。 - 它只返回那些至少存在一个在 k 近邻范围内的匹配邻居的对。 - 如果某个查询点没有有效邻居（例如 k 设得太大），它会被从结果中排除。 +- 空几何和 null 几何不会参与匹配。 ``` spark.sedona.join.knn.includeTieBreakers=true diff --git a/docs/tutorial/geopandas-api.md b/docs/tutorial/geopandas-api.md index 1c99c8711d5..eb953ef7000 100644 --- a/docs/tutorial/geopandas-api.md +++ b/docs/tutorial/geopandas-api.md @@ -237,6 +237,36 @@ intersects_result = left_df.sjoin(right_df, predicate="intersects") contains_result = left_df.sjoin(right_df, predicate="contains") ``` +### Nearest-Neighbor Joins + +`sjoin_nearest()` uses Sedona's distributed KNN join plan, so geometry rows are +not collected to the driver. It supports `inner`, `left`, and `right` results; +left and right joins retain unmatched rows from their respective side while +nearest matching remains distributed. + +```python +nearest = left_df.sjoin_nearest( + right_df, + how="left", + max_distance=500, + distance_col="distance", + exclusive=True, +) +``` + +Every candidate tied at the nearest distance is returned. `max_distance` +limits eligible matches, but it is currently applied after nearest candidates +are found and does not prune the distributed KNN search. `exclusive=True` +excludes candidates topologically equal to the query geometry before ranking. +Distances are planar and use CRS units; geographic CRS inputs emit a warning +because their distance results may be inaccurate. + +For inputs with MultiIndex columns or tuple-valued index names, the result uses +a padded MultiIndex when needed. This is the pandas-on-Spark representation of +GeoPandas' mixed tuple-and-string object column index, which pandas-on-Spark +cannot represent directly. The operation rejects suffixing or padding +combinations that would create duplicate labels. + ### Coordinate Reference System Operations Transform geometries between different coordinate reference systems: @@ -296,7 +326,7 @@ buildings_projected["perimeter"] = buildings_projected.geometry.length ## Supported Operations -The GeoPandas API for Apache Sedona has implemented **39 GeoSeries functions** and **10 GeoDataFrame functions**, covering the most commonly used GeoPandas operations: +The GeoPandas API for Apache Sedona implements the most commonly used GeoSeries and GeoDataFrame operations: ### Data I/O @@ -307,6 +337,7 @@ The GeoPandas API for Apache Sedona has implemented **39 GeoSeries functions** a ### Spatial Operations - `sjoin()` - Spatial joins with various predicates +- `sjoin_nearest()` - Distributed nearest-neighbor joins - `buffer()` - Geometric buffering - `distance()` - Distance calculations - `intersects()`, `contains()`, `within()` - Spatial predicates diff --git a/docs/tutorial/geopandas-api.zh.md b/docs/tutorial/geopandas-api.zh.md index fe6b3568613..558bd769d11 100644 --- a/docs/tutorial/geopandas-api.zh.md +++ b/docs/tutorial/geopandas-api.zh.md @@ -237,6 +237,32 @@ intersects_result = left_df.sjoin(right_df, predicate="intersects") contains_result = left_df.sjoin(right_df, predicate="contains") ``` +### 最近邻连接 + +`sjoin_nearest()` 使用 Sedona 的分布式 KNN 连接计划，不会把几何数据行收集到 +driver。它支持 `inner`、`left` 和 `right`；left 与 right 连接会保留对应一侧 +未匹配的行，而最近邻匹配本身仍以分布式方式执行。 + +```python +nearest = left_df.sjoin_nearest( + right_df, + how="left", + max_distance=500, + distance_col="distance", + exclusive=True, +) +``` + +所有在最近距离上并列的候选对象都会返回。`max_distance` 用于限制可匹配的 +距离，但目前是在找到最近候选对象后才应用，因此不会缩小分布式 KNN 搜索范围。 +`exclusive=True` 会在排序之前排除与查询几何拓扑相等的候选对象。距离按平面 +计算并使用 CRS 的单位；输入为地理坐标系时会发出警告，因为距离结果可能不准确。 + +对于使用 MultiIndex 列或元组型索引名称的输入，结果会在必要时使用补齐后的 +MultiIndex。这是 pandas-on-Spark 对 GeoPandas 混合元组与字符串对象列索引的表示 +方式；后者无法由 pandas-on-Spark 直接表达。如果后缀处理或补齐会产生重复标签， +该操作将拒绝执行。 + ### 坐标参考系操作 在不同坐标参考系（CRS）之间转换几何对象： @@ -296,7 +322,7 @@ buildings_projected["perimeter"] = buildings_projected.geometry.length ## 已支持的操作 -Apache Sedona 的 GeoPandas API 已实现 **39 个 GeoSeries 函数** 与 **10 个 GeoDataFrame 函数**，覆盖了 GeoPandas 中最常用的操作： +Apache Sedona 的 GeoPandas API 已实现最常用的 GeoSeries 与 GeoDataFrame 操作： ### 数据 I/O @@ -307,6 +333,7 @@ Apache Sedona 的 GeoPandas API 已实现 **39 个 GeoSeries 函数** 与 **10 ### 空间操作 - `sjoin()` —— 多种谓词的空间连接 +- `sjoin_nearest()` —— 分布式最近邻连接 - `buffer()` —— 几何缓冲 - `distance()` —— 距离计算 - `intersects()`、`contains()`、`within()` —— 空间谓词 diff --git a/python/sedona/spark/geopandas/__init__.py b/python/sedona/spark/geopandas/__init__.py index cf6c0af175d..788900fa917 100644 --- a/python/sedona/spark/geopandas/__init__.py +++ b/python/sedona/spark/geopandas/__init__.py @@ -23,7 +23,7 @@ from sedona.spark.geopandas.geoseries import GeoSeries from sedona.spark.geopandas.geodataframe import GeoDataFrame -from sedona.spark.geopandas.tools import sjoin +from sedona.spark.geopandas.tools import sjoin, sjoin_nearest from sedona.spark.geopandas.io import read_file, read_parquet diff --git a/python/sedona/spark/geopandas/geodataframe.py b/python/sedona/spark/geopandas/geodataframe.py index 68429c93e22..0a978728de1 100644 --- a/python/sedona/spark/geopandas/geodataframe.py +++ b/python/sedona/spark/geopandas/geodataframe.py @@ -33,6 +33,7 @@ from sedona.spark.geopandas._typing import Label from sedona.spark.geopandas.base import GeoFrame +from sedona.spark.sql import st_functions as stf from pandas.api.extensions import register_extension_dtype from geopandas.geodataframe import crs_mismatch_error @@ -338,7 +339,39 @@ def __init__( raise ValueError(crs_mismatch_error) if geometry: - self.set_geometry(geometry, inplace=True, crs=crs) + existing_geometry = None + if crs is not None and pd.api.types.is_hashable(geometry): + try: + candidate = self[geometry] + except (KeyError, TypeError, ValueError): + pass + else: + if isinstance(candidate, sgpd.GeoSeries): + existing_geometry = candidate + + if existing_geometry is None: + self.set_geometry(geometry, inplace=True, crs=crs) + else: + # Updating the existing Spark column directly avoids index + # alignment, which would multiply rows for duplicate indexes. + from pyproj import CRS + + normalized_crs = CRS.from_user_input(crs) + new_epsg = normalized_crs.to_epsg() or 0 + self._update_internal_frame( + self._internal.with_new_spark_column( + existing_geometry._column_label, + stf.ST_SetSRID(existing_geometry.spark.column, new_epsg), + ) + ) + self._geometry_column_name = geometry + empty_crs_source = self[geometry] + empty_crs_source._empty_crs_value = normalized_crs + object.__setattr__( + self, + "_empty_crs_source", + empty_crs_source, + ) if geometry is None and "geometry" in self.columns: @@ -396,7 +429,11 @@ def _get_geometry(self) -> sgpd.GeoSeries: ) raise MissingGeometryColumnError(msg) - return self[self._geometry_column_name] + geometry = self[self._geometry_column_name] + empty_crs_source = getattr(self, "_empty_crs_source", None) + if empty_crs_source is not None: + geometry._empty_crs_source = empty_crs_source + return geometry def _set_geometry(self, col): # This check is included in the original geopandas. Note that this prevents assigning a str to the property @@ -591,6 +628,7 @@ def set_geometry( if new_series: # Note: This casts GeoSeries back into pspd.Series, so we lose any metadata that's not serialized. frame[geo_column_name] = level + object.__setattr__(frame, "_empty_crs_source", level) if not inplace: return frame @@ -710,7 +748,14 @@ def _to_geopandas(self) -> gpd.GeoDataFrame: else: pd_df[col_name] = series._to_pandas() - return gpd.GeoDataFrame(pd_df, geometry=self._geometry_column_name) + result = gpd.GeoDataFrame( + pd_df, + geometry=self._geometry_column_name, + crs=self.crs if self._geometry_column_name is not None else None, + ) + if self._geometry_column_name is None: + result._geometry_column_name = None + return result def to_spark_pandas(self) -> pspd.DataFrame: """ @@ -1450,6 +1495,59 @@ def sjoin( **kwargs, ) + def sjoin_nearest( + self, + other, + how="inner", + max_distance=None, + lsuffix="left", + rsuffix="right", + distance_col=None, + exclusive=False, + ): + """Join another GeoDataFrame using distributed nearest neighbours. + + Parameters + ---------- + other : GeoDataFrame + Right-hand frame. + how : {"inner", "left", "right"}, default "inner" + Join mode and side whose index and active geometry are retained. + max_distance : float, optional + Maximum planar distance for a nearest match. This is applied after + nearest candidates are found and does not currently prune the KNN + search. + lsuffix, rsuffix : str + Suffixes for overlapping columns. + distance_col : str, optional + Name of a column containing match distances. + exclusive : bool, default False + Exclude topologically equal candidates. + + Returns + ------- + GeoDataFrame + Distributed nearest-join result. + + See Also + -------- + sedona.spark.geopandas.sjoin_nearest + """ + from sedona.spark.geopandas.tools.sjoin import ( + sjoin_nearest as sjoin_nearest_tool, + ) + + return sjoin_nearest_tool( + self, + other, + how=how, + max_distance=max_distance, + lsuffix=lsuffix, + rsuffix=rsuffix, + distance_col=distance_col, + exclusive=exclusive, + ) + # ============================================================================ # I/O OPERATIONS # ============================================================================ diff --git a/python/sedona/spark/geopandas/geoseries.py b/python/sedona/spark/geopandas/geoseries.py index 7ced034c51f..8f8ba97ed48 100644 --- a/python/sedona/spark/geopandas/geoseries.py +++ b/python/sedona/spark/geopandas/geoseries.py @@ -343,6 +343,9 @@ def __init__( self._col_label: Label self._sindex: SpatialIndex = None self._empty_crs_source: typing.Optional["GeoSeries"] = None + # Explicit CRS metadata wins over the lineage fallback below when + # geometry rows are empty or carry SRID 0. + self._empty_crs_value = None if isinstance( data, (GeoDataFrame, GeoSeries, PandasOnSparkSeries, PandasOnSparkDataFrame) @@ -461,6 +464,10 @@ def crs(self) -> Union["CRS", None]: from pyproj import CRS if self._is_empty(): + # Empty data has no SRID to inspect, so explicit CRS metadata wins + # over the inherited lineage metadata. + if self._empty_crs_value is not None: + return self._empty_crs_value if self._empty_crs_source is not None: return self._empty_crs_source.crs return None @@ -482,7 +489,16 @@ def crs(self) -> Union["CRS", None]: srid = 0 if np.isnan(srid) else srid # Sedona returns 0 if SRID doesn't exist. - return CRS.from_user_input(srid) if srid != 0 else None + if srid != 0: + return CRS.from_user_input(srid) + # These fallbacks are metadata rather than a fresh read from geometry + # coordinates. Explicit metadata takes precedence over inherited + # lineage metadata, including for non-empty geometries with SRID 0. + if self._empty_crs_value is not None: + return self._empty_crs_value + if self._empty_crs_source is not None: + return self._empty_crs_source.crs + return None @crs.setter def crs(self, value: Union["CRS", None]): @@ -615,9 +631,15 @@ def set_crs( spark_col = stf.ST_SetSRID(self.spark.column, new_epsg) result = self._query_geometry_column(spark_col, keep_name=True) + result._empty_crs_value = crs + if crs is None: + result._empty_crs_source = None if inplace: self._update_inplace(result, invalidate_sindex=False) + self._empty_crs_value = crs + if crs is None: + self._empty_crs_source = None return None return result diff --git a/python/sedona/spark/geopandas/tools/__init__.py b/python/sedona/spark/geopandas/tools/__init__.py index 05067bba3de..f3435f62023 100644 --- a/python/sedona/spark/geopandas/tools/__init__.py +++ b/python/sedona/spark/geopandas/tools/__init__.py @@ -15,8 +15,9 @@ # specific language governing permissions and limitations # under the License. -from .sjoin import sjoin +from .sjoin import sjoin, sjoin_nearest __all__ = [ "sjoin", + "sjoin_nearest", ] diff --git a/python/sedona/spark/geopandas/tools/sjoin.py b/python/sedona/spark/geopandas/tools/sjoin.py index 057190ade5b..fa02c717be8 100644 --- a/python/sedona/spark/geopandas/tools/sjoin.py +++ b/python/sedona/spark/geopandas/tools/sjoin.py @@ -14,18 +14,503 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import numbers import re +import warnings +from dataclasses import dataclass +from typing import Any, List, Optional, Tuple + +import numpy as np import pyspark.pandas as ps -from pyspark.pandas.internal import SPARK_DEFAULT_INDEX_NAME, InternalFrame +from pyspark.pandas.internal import ( + NATURAL_ORDER_COLUMN_NAME, + SPARK_DEFAULT_INDEX_NAME, + InternalFrame, +) from pyspark.pandas.utils import scol_for +from pyspark.sql import Column, DataFrame as SparkDataFrame +from pyspark.sql import functions as F from pyspark.sql.functions import expr from sedona.spark.geopandas import GeoDataFrame, GeoSeries +from sedona.spark.sql import st_functions as stf +from sedona.spark.sql import st_predicates as stp # Pre-compiled regex pattern for suffix validation. SUFFIX_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") +@dataclass(frozen=True) +class _NearestJoinSide: + """A join input projected onto collision-proof Spark column names.""" + + spark_frame: SparkDataFrame + index_aliases: Tuple[str, ...] + index_names: Tuple[Any, ...] + data_aliases: Tuple[str, ...] + data_labels: Tuple[Tuple[Any, ...], ...] + active_label: Tuple[Any, ...] + geometry_alias: str + order_alias: str + + +def _qualified_column(qualifier: str, name: str) -> Column: + return F.col(f"{qualifier}.`{name}`") + + +def _nearest_label_in_columns( + label: Tuple[Any, ...], + column_labels: List[Tuple[Any, ...]], +) -> bool: + """Apply pandas Index/MultiIndex membership rules to a column label.""" + + if label in column_labels: + return True + return len(label) == 1 and any( + len(column_label) > 1 and column_label[0] == label[0] + for column_label in column_labels + ) + + +def _prepare_nearest_side(frame: GeoDataFrame, prefix: str) -> _NearestJoinSide: + internal = frame._internal + active_label = frame.geometry._column_label + if active_label is None: + raise ValueError("GeoDataFrame geometry column is not set") + + try: + geometry_position = internal.column_labels.index(active_label) + except ValueError as exc: + raise ValueError("GeoDataFrame geometry column is not present") from exc + + index_aliases = tuple( + f"__sjoin_nearest_{prefix}_index_{position}__" + for position in range(len(internal.index_spark_columns)) + ) + data_aliases = tuple( + f"__sjoin_nearest_{prefix}_data_{position}__" + for position in range(len(internal.data_spark_columns)) + ) + order_alias = f"__sjoin_nearest_{prefix}_order__" + + projected = internal.spark_frame.select( + *[ + column.alias(alias) + for column, alias in zip(internal.index_spark_columns, index_aliases) + ], + *[ + column.alias(alias) + for column, alias in zip(internal.data_spark_columns, data_aliases) + ], + scol_for(internal.spark_frame, NATURAL_ORDER_COLUMN_NAME).alias(order_alias), + ) + return _NearestJoinSide( + spark_frame=projected, + index_aliases=index_aliases, + index_names=tuple(internal.index_names), + data_aliases=data_aliases, + data_labels=tuple(internal.column_labels), + active_label=active_label, + geometry_alias=data_aliases[geometry_position], + order_alias=order_alias, + ) + + +def _nearest_reset_index_labels( + side: _NearestJoinSide, + suffix: str, + side_data_labels: List[Tuple[Any, ...]], + other_data_labels: List[Tuple[Any, ...]], +) -> List[Tuple[Any, ...]]: + """Return GeoPandas-compatible labels for reset index levels.""" + + labels: List[Tuple[Any, ...]] = [] + multiple_levels = len(side.index_names) > 1 + for position, name in enumerate(side.index_names): + if name is None: + generated = ( + f"index_{suffix}{position}" if multiple_levels else f"index_{suffix}" + ) + label = (generated,) + if _nearest_label_in_columns( + label, side_data_labels + ) or _nearest_label_in_columns(label, other_data_labels): + raise ValueError( + f"'{generated}' cannot be a column name in the frames being joined" + ) + else: + label = name + if _nearest_label_in_columns(label, side_data_labels): + display_name = label[0] if len(label) == 1 else label + raise ValueError(f"cannot insert {display_name}, already exists") + labels.append(label) + return labels + + +def _nearest_suffix_label(label: Tuple[Any, ...], suffix: str) -> Tuple[Any, ...]: + if len(label) == 1: + return (f"{label[0]}_{suffix}",) + return (f"{label}_{suffix}",) + + +def _nearest_output_columns( + left: _NearestJoinSide, + right: _NearestJoinSide, + how: str, + lsuffix: str, + rsuffix: str, +) -> List[Tuple[str, Tuple[Any, ...]]]: + """Resolve output values and labels using GeoPandas reset/suffix rules.""" + + left_data = [ + (alias, label) + for alias, label in zip(left.data_aliases, left.data_labels) + if how != "right" or alias != left.geometry_alias + ] + right_data = [ + (alias, label) + for alias, label in zip(right.data_aliases, right.data_labels) + if how == "right" or alias != right.geometry_alias + ] + left_data_labels = [label for _, label in left_data] + right_data_labels = [label for _, label in right_data] + left_index_labels = _nearest_reset_index_labels( + left, + lsuffix, + left_data_labels, + right_data_labels, + ) + right_index_labels = _nearest_reset_index_labels( + right, + rsuffix, + right_data_labels, + left_data_labels, + ) + + left_reset_labels = left_index_labels + left_data_labels + right_reset_labels = right_index_labels + right_data_labels + overlap = set(left_reset_labels).intersection(right_reset_labels) + + def rename( + label: Tuple[Any, ...], + suffix: str, + geometry_label: Tuple[Any, ...], + ) -> Tuple[Any, ...]: + if label in overlap and label != geometry_label: + return _nearest_suffix_label(label, suffix) + return label + + output: List[Tuple[str, Tuple[Any, ...]]] = [] + if how == "right": + output.extend( + (alias, rename(label, lsuffix, left.active_label)) + for alias, label in zip(left.index_aliases, left_index_labels) + ) + output.extend( + (alias, rename(label, lsuffix, left.active_label)) for alias, label in left_data + ) + + if how in ("inner", "left"): + output.extend( + (alias, rename(label, rsuffix, right.active_label)) + for alias, label in zip(right.index_aliases, right_index_labels) + ) + output.extend( + (alias, rename(label, rsuffix, right.active_label)) + for alias, label in right_data + ) + return output + + +def _nearest_result_index_names( + left: _NearestJoinSide, + right: _NearestJoinSide, + how: str, + lsuffix: str, + rsuffix: str, +) -> List[Any]: + """Apply the reset/merge suffix rules to the retained index names.""" + left_data_labels = [ + label + for alias, label in zip(left.data_aliases, left.data_labels) + if how != "right" or alias != left.geometry_alias + ] + right_data_labels = [ + label + for alias, label in zip(right.data_aliases, right.data_labels) + if how == "right" or alias != right.geometry_alias + ] + left_index_labels = _nearest_reset_index_labels( + left, + lsuffix, + left_data_labels, + right_data_labels, + ) + right_index_labels = _nearest_reset_index_labels( + right, + rsuffix, + right_data_labels, + left_data_labels, + ) + overlap = set(left_index_labels + left_data_labels).intersection( + right_index_labels + right_data_labels + ) + + if how in ("inner", "left"): + side, labels, suffix = left, left_index_labels, lsuffix + else: + side, labels, suffix = right, right_index_labels, rsuffix + + return [ + ( + _nearest_suffix_label(label, suffix) + if name is not None and label in overlap + else name + ) + for name, label in zip(side.index_names, labels) + ] + + +def _normalize_nearest_output_labels( + output_columns: List[Tuple[str, Tuple[Any, ...]]], +) -> Tuple[List[Tuple[str, Tuple[Any, ...]]], int]: + """Pad labels to the uniform depth required by pandas-on-Spark.""" + + output_levels = max((len(label) for _, label in output_columns), default=1) + if output_levels == 1: + return output_columns, output_levels + return ( + [ + (alias, label + ("",) * (output_levels - len(label))) + for alias, label in output_columns + ], + output_levels, + ) + + +def _validate_nearest_output_labels( + output_columns: List[Tuple[str, Tuple[Any, ...]]], +) -> None: + """Reject duplicate labels that pandas-on-Spark cannot represent safely.""" + + seen_labels = set() + duplicate_labels = set() + for _, label in output_columns: + if label in seen_labels: + duplicate_labels.add(label) + seen_labels.add(label) + if duplicate_labels: + display_labels = ", ".join( + sorted( + repr(label[0] if len(label) == 1 else label) + for label in duplicate_labels + ) + ) + raise ValueError( + "sjoin_nearest cannot represent duplicate output columns after " + f"suffixing or MultiIndex padding: {display_labels}" + ) + + +def _check_sjoin_nearest_crs(left_df: GeoDataFrame, right_df: GeoDataFrame) -> None: + left_crs = left_df.crs + right_crs = right_df.crs + + if left_crs != right_crs: + left_crs_text = "None" if left_crs is None else left_crs.to_string() + right_crs_text = "None" if right_crs is None else right_crs.to_string() + warnings.warn( + "CRS mismatch between the CRS of left geometries and the CRS of " + "right geometries.\nUse `to_crs()` to reproject one of the input " + "geometries to match the CRS of the other.\n\n" + f"Left CRS: {left_crs_text}\n" + f"Right CRS: {right_crs_text}\n", + UserWarning, + stacklevel=3, + ) + + geographic_warning = ( + "Geometry is in a geographic CRS. Results from 'sjoin_nearest' are " + "likely incorrect. Use 'GeoSeries.to_crs()' to re-project geometries " + "to a projected CRS before this operation.\n" + ) + for crs in (left_crs, right_crs): + if crs is not None and crs.is_geographic: + warnings.warn( + geographic_warning, + UserWarning, + stacklevel=3, + ) + + +def _normalize_max_distance(max_distance) -> Optional[float]: + if max_distance is None: + return None + if not isinstance(max_distance, numbers.Real): + raise TypeError("max_distance must be a number") + max_distance = float(max_distance) + if max_distance <= 0: + raise ValueError("max_distance must be greater than 0") + return max_distance + + +def _nearest_frame_join( + left_df: GeoDataFrame, + right_df: GeoDataFrame, + how: str, + max_distance: Optional[float], + lsuffix: str, + rsuffix: str, + distance_col: Optional[str], + exclusive: bool, +) -> GeoDataFrame: + left = _prepare_nearest_side(left_df, "left") + right = _prepare_nearest_side(right_df, "right") + + if how == "right": + query, objects = right, left + left_qualifier, right_qualifier = "o", "q" + else: + query, objects = left, right + left_qualifier, right_qualifier = "q", "o" + + query_frame = query.spark_frame.alias("q") + object_frame = objects.spark_frame.alias("o") + query_geometry = _qualified_column("q", query.geometry_alias) + object_geometry = _qualified_column("o", objects.geometry_alias) + distance = stf.ST_Distance(query_geometry, object_geometry) + condition = stp.ST_KNN( + query_geometry, + object_geometry, + 1, + use_spheroid=False, + include_ties=True, + exclusive=exclusive, + ) + if max_distance is not None: + condition = condition & (distance <= F.lit(max_distance)) + + matches = query_frame.join(object_frame, condition, "inner") + distance_alias = "__sjoin_nearest_distance__" + + if how == "inner": + joined = matches.select( + *[ + _qualified_column(left_qualifier, name).alias(name) + for name in left.spark_frame.columns + ], + *[ + _qualified_column(right_qualifier, name).alias(name) + for name in right.spark_frame.columns + ], + *([distance.alias(distance_alias)] if distance_col is not None else []), + ) + else: + match_id = "__sjoin_nearest_match_id__" + payload = matches.select( + _qualified_column("q", query.order_alias).alias(match_id), + *[ + _qualified_column("o", name).alias(name) + for name in objects.spark_frame.columns + ], + *([distance.alias(distance_alias)] if distance_col is not None else []), + ) + joined = query.spark_frame.alias("base").join( + payload.alias("match"), + _qualified_column("base", query.order_alias) + == _qualified_column("match", match_id), + "left", + ) + + active = left if how in ("inner", "left") else right + opposite = right if how in ("inner", "left") else left + joined = joined.orderBy( + F.col(active.order_alias).asc(), + F.col(opposite.order_alias).asc_nulls_last(), + ) + + output_columns = _nearest_output_columns( + left, + right, + how, + lsuffix, + rsuffix, + ) + distance_replaces_geometry = False + if distance_col is not None: + distance_label = (distance_col,) + distance_replaces_geometry = distance_label == active.active_label + replaced = False + for position, (_, label) in enumerate(output_columns): + if label == distance_label: + output_columns[position] = (distance_alias, distance_label) + replaced = True + break + if not replaced: + output_columns.append((distance_alias, distance_label)) + output_columns, output_column_levels = _normalize_nearest_output_labels( + output_columns + ) + _validate_nearest_output_labels(output_columns) + output_column_label_names = [None] * output_column_levels + + index_names = [ + f"__sjoin_nearest_result_index_{position}__" + for position in range(len(active.index_aliases)) + ] + data_names = [ + f"__sjoin_nearest_result_data_{position}__" + for position in range(len(output_columns)) + ] + result_frame = joined.select( + *[ + F.col(alias).alias(name) + for alias, name in zip(active.index_aliases, index_names) + ], + *[ + F.col(alias).alias(name) + for (alias, _), name in zip(output_columns, data_names) + ], + ) + + internal = InternalFrame( + spark_frame=result_frame, + index_spark_columns=[scol_for(result_frame, name) for name in index_names], + index_names=_nearest_result_index_names( + left, + right, + how, + lsuffix, + rsuffix, + ), + column_labels=[label for _, label in output_columns], + data_spark_columns=[scol_for(result_frame, name) for name in data_names], + column_label_names=output_column_label_names, + ) + result = GeoDataFrame(ps.DataFrame(internal)) + if distance_replaces_geometry: + warnings.warn( + "Geometry column does not contain geometry.", + UserWarning, + stacklevel=3, + ) + object.__setattr__(result, "_geometry_column_name", None) + else: + active_output_label = active.active_label + ("",) * ( + output_column_levels - len(active.active_label) + ) + active_output_name = ( + active_output_label[0] if output_column_levels == 1 else active_output_label + ) + object.__setattr__(result, "_geometry_column_name", active_output_name) + object.__setattr__( + result, + "_empty_crs_source", + left_df.geometry if how in ("inner", "left") else right_df.geometry, + ) + return result + + def _frame_join( left_df: GeoDataFrame, right_df: GeoDataFrame, @@ -314,6 +799,110 @@ def sjoin( return joined +def sjoin_nearest( + left_df: GeoDataFrame, + right_df: GeoDataFrame, + how: str = "inner", + max_distance: Optional[float] = None, + lsuffix: str = "left", + rsuffix: str = "right", + distance_col: Optional[str] = None, + exclusive: bool = False, +) -> GeoDataFrame: + """Spatial join based on the nearest distributed geometry matches. + + Every equidistant nearest match is returned. The operation is planned as + Sedona's distributed KNN join and does not collect geometry rows to the + driver. + + Parameters + ---------- + left_df, right_df : GeoDataFrame + Frames to join. + how : {"inner", "left", "right"}, default "inner" + Join mode and the side whose index and active geometry are retained. + max_distance : float, optional + Maximum planar distance for a nearest match. Must be greater than zero. + This is currently applied after nearest candidates are found, so it + limits output matches but does not prune the distributed KNN search. + lsuffix, rsuffix : str + Suffixes applied to overlapping left and right column names. + distance_col : str, optional + Output column containing the planar distance between each match. + exclusive : bool, default False + Exclude candidates topologically equal to the query geometry before + nearest matches are ranked. + + Returns + ------- + GeoDataFrame + Distributed nearest-join result. + + Examples + -------- + >>> from shapely.geometry import Point + >>> from sedona.spark.geopandas import GeoDataFrame, sjoin_nearest + >>> left = GeoDataFrame( + ... {"name": ["a"], "geometry": [Point(0, 0)]}, crs="EPSG:3857" + ... ) + >>> right = GeoDataFrame( + ... {"value": [1, 2], "geometry": [Point(-1, 0), Point(1, 0)]}, + ... crs="EPSG:3857", + ... ) + >>> sjoin_nearest(left, right, distance_col="distance").sort_values("value") + name geometry index_right value distance + 0 a POINT (0 0) 0 1 1.0 + 0 a POINT (0 0) 1 2 1.0 + + Notes + ----- + Distances are planar and expressed in CRS units. Geographic CRS inputs + therefore emit the same accuracy warning as GeoPandas. + + ``max_distance`` is currently a post-filter rather than an index search + bound. It preserves result semantics but does not reduce KNN search work. + + When either input has MultiIndex columns or tuple-valued index names, + pandas-on-Spark cannot represent GeoPandas' one-level object index + containing both tuple and string labels. Sedona therefore returns an + equivalent padded MultiIndex when needed and preserves the retained active + geometry label. Outputs whose suffixing or padding would collapse distinct + labels to duplicates are rejected because pandas-on-Spark cannot represent + them safely. + """ + _basic_checks(left_df, right_df, how, lsuffix, rsuffix) + + if left_df.active_geometry_name is None: + raise ValueError("Left DataFrame geometry column not set") + if right_df.active_geometry_name is None: + raise ValueError("Right DataFrame geometry column not set") + if isinstance(exclusive, np.bool_): + exclusive = bool(exclusive) + elif isinstance(exclusive, numbers.Integral): + if exclusive not in (0, 1): + raise ValueError("exclusive parameter must be boolean") + exclusive = bool(exclusive) + elif not isinstance(exclusive, bool): + if not np.isscalar(exclusive): + raise ValueError("exclusive parameter only accepts scalar values") + raise ValueError("exclusive parameter must be boolean") + if distance_col is not None and not isinstance(distance_col, str): + raise TypeError("distance_col must be a string or None") + + max_distance = _normalize_max_distance(max_distance) + _check_sjoin_nearest_crs(left_df, right_df) + return _nearest_frame_join( + left_df, + right_df, + how, + max_distance, + lsuffix, + rsuffix, + distance_col, + exclusive, + ) + + def _maybe_make_list(obj): if isinstance(obj, tuple): return list(obj) diff --git a/python/sedona/spark/sql/st_predicates.py b/python/sedona/spark/sql/st_predicates.py index dd04b77d1c8..e2ef797eb3c 100644 --- a/python/sedona/spark/sql/st_predicates.py +++ b/python/sedona/spark/sql/st_predicates.py @@ -133,6 +133,54 @@ def ST_Intersects(a: ColumnOrName, b: ColumnOrName) -> Column: return _call_predicate_function("ST_Intersects", (a, b)) +@validate_argument_types +def ST_KNN( + a: ColumnOrName, + b: ColumnOrName, + k: Union[ColumnOrName, int], + use_spheroid: Optional[Union[ColumnOrName, bool]] = None, + include_ties: Optional[Union[ColumnOrName, bool]] = None, + exclusive: Optional[Union[ColumnOrName, bool]] = None, +) -> Column: + """Build an optimizer-recognized K-nearest-neighbour join predicate. + + ``include_ties`` and ``exclusive`` are query-local controls. When + ``include_ties`` is omitted, the existing + ``spark.sedona.join.knn.includeTieBreakers`` session setting is used. + ``exclusive=True`` excludes topologically equal candidates before nearest + neighbours are ranked. + + This predicate is only valid as an inner join condition. + + :param a: Query geometry column. + :param b: Candidate geometry column. + :param k: Number of nearest neighbours. + :param use_spheroid: Whether to use spheroid distance. + :param include_ties: Whether to return every candidate tied at the kth distance. + :param exclusive: Whether to exclude candidates equal to the query geometry. + :return: A KNN join predicate column. + """ + if exclusive is not None and include_ties is None: + raise ValueError("exclusive requires include_ties to be specified") + + if include_ties is not None: + args = ( + a, + b, + k, + False if use_spheroid is None else use_spheroid, + include_ties, + ) + if exclusive is not None: + args += (exclusive,) + elif use_spheroid is not None: + args = (a, b, k, use_spheroid) + else: + args = (a, b, k) + + return _call_predicate_function("ST_KNN", args) + + @validate_argument_types def ST_OrderingEquals(a: ColumnOrName, b: ColumnOrName) -> Column: """Check whether two geometries have identical vertices that are in the same order. diff --git a/python/tests/geopandas/test_sjoin_nearest.py b/python/tests/geopandas/test_sjoin_nearest.py new file mode 100644 index 00000000000..a26be476c35 --- /dev/null +++ b/python/tests/geopandas/test_sjoin_nearest.py @@ -0,0 +1,635 @@ +# 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 geopandas as gpd +import numpy as np +import pandas as pd +import pytest +import pyspark.pandas as ps +import shapely +from geopandas.testing import assert_geodataframe_equal +from packaging.version import parse as parse_version +from shapely.geometry import LineString, Point + +from sedona.spark.geopandas import GeoDataFrame, sjoin_nearest +from sedona.spark.sql import st_functions as stf +from tests.geopandas.test_geopandas_base import TestGeopandasBase + +pytestmark = pytest.mark.skipif( + parse_version(shapely.__version__) < parse_version("2.0.0"), + reason=f"Tests require shapely>=2.0.0, but found v{shapely.__version__}", +) + + +class TestSpatialJoinNearest(TestGeopandasBase): + @staticmethod + def _to_sedona(frame): + with ps.option_context("compute.ops_on_diff_frames", True): + return GeoDataFrame(frame) + + @staticmethod + def _frames(): + left = gpd.GeoDataFrame( + { + "key": ["a", "b"], + "same": [1, 2], + "geometry": [Point(0, 0), Point(10, 0)], + }, + index=pd.Index([5, 6], name="left_id"), + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + { + "same": [3, 4, 5], + "value": ["west", "east", "far"], + "geometry": [Point(-1, 0), Point(1, 0), Point(30, 0)], + }, + index=pd.Index([7, 8, 9], name="right_id"), + crs="EPSG:3857", + ) + return left, right + + @staticmethod + def _assert_matches_geopandas(actual, expected): + assert isinstance(actual, GeoDataFrame) + assert_geodataframe_equal( + actual.to_geopandas(), + expected, + check_dtype=False, + check_index_type=False, + ) + + @pytest.mark.parametrize("how", ["inner", "left", "right"]) + def test_ties_join_modes_distance_and_method(self, how): + left, right = self._frames() + expected = left.sjoin_nearest( + right, + how=how, + distance_col="distance", + ) + left_sgpd = self._to_sedona(left) + right_sgpd = self._to_sedona(right) + + actual = sjoin_nearest( + left_sgpd, + right_sgpd, + how=how, + distance_col="distance", + ) + method_result = left_sgpd.sjoin_nearest( + right_sgpd, + how=how, + distance_col="distance", + ) + + self._assert_matches_geopandas(actual, expected) + self._assert_matches_geopandas(method_result, expected) + assert actual.crs == expected.crs + + @pytest.mark.parametrize("how", ["inner", "left", "right"]) + def test_max_distance(self, how): + left, right = self._frames() + expected = left.sjoin_nearest( + right, + how=how, + max_distance=2.0, + distance_col="distance", + ) + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + max_distance=2.0, + distance_col="distance", + ) + self._assert_matches_geopandas(actual, expected) + + @pytest.mark.parametrize("how", ["inner", "right"]) + def test_duplicate_multiindex_and_index_name_suffixes(self, how): + left_index = pd.MultiIndex.from_tuples( + [(1, "a"), (1, "a")], + names=["id", "part"], + ) + right_index = pd.MultiIndex.from_tuples( + [(2, "b")], + names=["id", "part"], + ) + left = gpd.GeoDataFrame( + { + "shared": [1, 2], + "geometry": [Point(0, 0), Point(0, 0)], + }, + index=left_index, + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + {"shared": [3], "geometry": [Point(1, 0)]}, + index=right_index, + crs="EPSG:3857", + ) + + expected = left.sjoin_nearest( + right, + how=how, + distance_col="distance", + ) + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + distance_col="distance", + ) + self._assert_matches_geopandas(actual, expected) + assert actual.index.names == expected.index.names + + @pytest.mark.parametrize("how", ["inner", "left", "right"]) + def test_multiindex_columns_use_padded_output(self, how): + columns = pd.MultiIndex.from_tuples( + [("a", "value"), ("g", "geometry")], + names=["kind", "field"], + ) + left = gpd.GeoDataFrame( + [[1, Point(0, 0)]], + columns=columns, + geometry=("g", "geometry"), + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + [[2, Point(1, 0)]], + columns=columns, + geometry=("g", "geometry"), + crs="EPSG:3857", + ) + distance_col = "('a', 'value')_left" if how == "inner" else "distance" + + expected = left.sjoin_nearest( + right, + how=how, + distance_col=distance_col, + ) + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + distance_col=distance_col, + ) + collected = actual.to_geopandas() + + assert isinstance(actual.columns, pd.MultiIndex) + assert actual.columns.names == [None, None] + expected_padded_columns = pd.MultiIndex.from_tuples( + [ + label if isinstance(label, tuple) else (label, "") + for label in expected.columns + ] + ) + assert actual.columns.equals(expected_padded_columns) + expected_active = ("g", "geometry") + assert actual.active_geometry_name == expected_active + assert actual.crs == expected.crs + + # GeoPandas flattens reset-index and suffixed labels into a mixed + # one-level object Index. pandas-on-Spark cannot represent that Index, + # so compare values after assigning the oracle's labels. + collected.columns = expected.columns + collected._geometry_column_name = expected.active_geometry_name + assert_geodataframe_equal( + collected, + expected, + check_dtype=False, + check_index_type=False, + ) + + @pytest.mark.parametrize("how", ["inner", "left", "right"]) + def test_column_axis_name_is_dropped(self, how): + left = gpd.GeoDataFrame( + {"a": [1], "geometry": [Point(0, 0)]}, + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + {"b": [2], "geometry": [Point(1, 0)]}, + crs="EPSG:3857", + ) + left.columns.name = "left_columns" + right.columns.name = "right_columns" + + expected = left.sjoin_nearest(right, how=how) + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + ) + + self._assert_matches_geopandas(actual, expected) + assert actual.columns.name is None + + @pytest.mark.parametrize("how", ["inner", "right"]) + def test_tuple_index_name_suffixes_match_geopandas(self, how): + index = pd.Index([0], name=("idx", "part")) + left = gpd.GeoDataFrame( + {"a": [1], "geometry": [Point(0, 0)]}, + index=index, + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + {"b": [2], "geometry": [Point(1, 0)]}, + index=index, + crs="EPSG:3857", + ) + + expected = left.sjoin_nearest(right, how=how) + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + ) + + self._assert_matches_geopandas(actual, expected) + assert list(actual.columns) == list(expected.columns) + assert actual.index.names == expected.index.names + + def test_suffix_created_duplicate_labels_are_rejected(self): + left = gpd.GeoDataFrame( + { + "a": [10], + "a_left": [20], + "geometry": [Point(0, 0)], + }, + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + { + "a": [30], + "geometry": [Point(3, 4)], + }, + crs="EPSG:3857", + ) + + with pytest.warns(FutureWarning): + left.sjoin_nearest( + right, + distance_col="a_left", + ) + with pytest.raises(ValueError, match="duplicate output columns.*a_left"): + sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + distance_col="a_left", + ) + + def test_multiindex_padding_collision_is_rejected(self): + left = gpd.GeoDataFrame( + {"a": [1], "geometry": [Point(0, 0)]}, + crs="EPSG:3857", + ) + right_columns = pd.MultiIndex.from_tuples([("a", ""), ("g", "geometry")]) + right = gpd.GeoDataFrame( + [[2, Point(1, 0)]], + columns=right_columns, + geometry=("g", "geometry"), + crs="EPSG:3857", + ) + + expected = left.sjoin_nearest(right) + assert expected.columns.is_unique + with pytest.raises(ValueError, match="duplicate output columns.*a"): + sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + ) + + @pytest.mark.parametrize( + "multi_on_left,how", + [(True, "right"), (False, "inner")], + ) + def test_mixed_column_depth_preserves_active_geometry( + self, + multi_on_left, + how, + ): + columns = pd.MultiIndex.from_tuples( + [("a", "value"), ("g", "geometry")], + names=["kind", "field"], + ) + multi = gpd.GeoDataFrame( + [[1, Point(0, 0)]], + columns=columns, + geometry=("g", "geometry"), + crs="EPSG:3857", + ) + plain = gpd.GeoDataFrame( + {"a": [2], "geometry": [Point(1, 0)]}, + crs="EPSG:3857", + ) + left, right = (multi, plain) if multi_on_left else (plain, multi) + + expected = left.sjoin_nearest(right, how=how) + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + ) + collected = actual.to_geopandas() + + assert isinstance(actual.columns, pd.MultiIndex) + assert actual.active_geometry_name == ("geometry", "") + assert actual.crs == expected.crs + collected.columns = expected.columns + collected._geometry_column_name = expected.active_geometry_name + assert_geodataframe_equal( + collected, + expected, + check_dtype=False, + check_index_type=False, + ) + + @pytest.mark.parametrize("how", ["inner", "left"]) + def test_dropped_right_geometry_does_not_conflict_with_left_index(self, how): + left = gpd.GeoDataFrame( + {"geometry": [Point(0, 0)]}, + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + { + "x": [1], + "index_left": [Point(1, 0)], + }, + geometry="index_left", + crs="EPSG:3857", + ) + + expected = left.sjoin_nearest(right, how=how) + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + ) + self._assert_matches_geopandas(actual, expected) + + def test_dropped_left_geometry_does_not_conflict_with_right_index(self): + left = gpd.GeoDataFrame( + { + "x": [1], + "index_right": [Point(0, 0)], + }, + geometry="index_right", + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + {"geometry": [Point(1, 0)]}, + crs="EPSG:3857", + ) + + expected = left.sjoin_nearest(right, how="right") + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how="right", + ) + self._assert_matches_geopandas(actual, expected) + + @pytest.mark.parametrize( + "column_name,how", + [("index_left", "inner"), ("index_right", "right")], + ) + def test_generated_index_detects_multiindex_first_level_collision( + self, + column_name, + how, + ): + columns = pd.MultiIndex.from_tuples([(column_name, "value"), ("g", "geometry")]) + conflicting = gpd.GeoDataFrame( + [[1, Point(0, 0)]], + columns=columns, + geometry=("g", "geometry"), + crs="EPSG:3857", + ) + plain = gpd.GeoDataFrame( + {"geometry": [Point(1, 0)]}, + crs="EPSG:3857", + ) + left, right = ( + (conflicting, plain) + if column_name == "index_left" + else (plain, conflicting) + ) + + with pytest.raises(ValueError, match=column_name): + left.sjoin_nearest(right, how=how) + with pytest.raises(ValueError, match=column_name): + sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + ) + + def test_exclusive_uses_topological_equality(self): + left = gpd.GeoDataFrame( + {"geometry": [LineString([(0, 0), (1, 0)])]}, + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + { + "value": ["reversed", "next"], + "geometry": [ + LineString([(1, 0), (0, 0)]), + LineString([(0, 2), (1, 2)]), + ], + }, + crs="EPSG:3857", + ) + + expected = left.sjoin_nearest( + right, + exclusive=True, + distance_col="distance", + ) + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + exclusive=True, + distance_col="distance", + ) + self._assert_matches_geopandas(actual, expected) + + def test_exclusive_self_join(self): + frame = gpd.GeoDataFrame( + { + "value": ["a", "b", "c"], + "geometry": [Point(0, 0), Point(2, 0), Point(5, 0)], + }, + crs="EPSG:3857", + ) + expected = frame.sjoin_nearest( + frame, + exclusive=True, + distance_col="distance", + ) + sedona_frame = self._to_sedona(frame) + actual = sedona_frame.sjoin_nearest( + sedona_frame, + exclusive=True, + distance_col="distance", + ) + self._assert_matches_geopandas(actual, expected) + + @pytest.mark.parametrize("how", ["inner", "left", "right"]) + def test_null_and_empty_geometries(self, how): + left = gpd.GeoDataFrame( + { + "value": ["point", "null", "empty"], + "geometry": [Point(0, 0), None, Point()], + }, + crs="EPSG:3857", + ) + right = gpd.GeoDataFrame( + {"candidate": [1], "geometry": [Point(1, 0)]}, + crs="EPSG:3857", + ) + + expected = left.sjoin_nearest(right, how=how, distance_col="distance") + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + distance_col="distance", + ) + self._assert_matches_geopandas(actual, expected) + + @pytest.mark.parametrize("all_null", [False, True]) + def test_empty_and_all_null_candidate_frames(self, all_null): + left = gpd.GeoDataFrame( + {"value": [1], "geometry": [Point(0, 0)]}, + crs="EPSG:3857", + ) + if all_null: + right = gpd.GeoDataFrame( + {"candidate": [1, 2], "geometry": [None, None]}, + crs="EPSG:3857", + ) + else: + right = gpd.GeoDataFrame( + { + "candidate": pd.Series(dtype="int64"), + "geometry": gpd.GeoSeries([], crs="EPSG:3857"), + }, + crs="EPSG:3857", + ) + + for how in ("inner", "left", "right"): + expected = left.sjoin_nearest( + right, + how=how, + distance_col="distance", + ) + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + how=how, + distance_col="distance", + ) + self._assert_matches_geopandas(actual, expected) + assert actual.crs == expected.crs + + def test_custom_geometry_name_crs_srid_and_plan(self): + left, right = self._frames() + left = left.rename_geometry("left_geometry") + right = right.rename_geometry("right_geometry") + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + ) + + assert actual.active_geometry_name == "left_geometry" + assert actual.crs == left.crs + srids = actual._internal.spark_frame.select( + stf.ST_SRID(actual.geometry.spark.column).alias("srid") + ).collect() + assert {row.srid for row in srids if row.srid is not None} == {3857} + + plan = ( + actual._internal.spark_frame._jdf.queryExecution().executedPlan().toString() + ) + assert "KNNJoin" in plan + assert "CartesianProduct" not in plan + assert "BatchEvalPython" not in plan + assert "PythonUDF" not in plan + + def test_distance_column_can_replace_active_geometry(self): + left, right = self._frames() + expected = left.sjoin_nearest(right, distance_col="geometry") + with pytest.warns(UserWarning, match="does not contain geometry"): + actual = sjoin_nearest( + self._to_sedona(left), + self._to_sedona(right), + distance_col="geometry", + ) + + assert actual.active_geometry_name is None + pd.testing.assert_frame_equal( + pd.DataFrame(actual.to_geopandas()), + pd.DataFrame(expected), + check_dtype=False, + check_index_type=False, + ) + + def test_validation_and_crs_warnings(self): + left, right = self._frames() + left_sgpd = self._to_sedona(left) + right_sgpd = self._to_sedona(right) + + with pytest.raises(ValueError, match="expected to be in"): + sjoin_nearest(left_sgpd, right_sgpd, how="outer") + with pytest.raises(ValueError, match="greater than 0"): + sjoin_nearest(left_sgpd, right_sgpd, max_distance=0) + with pytest.raises(TypeError, match="must be a number"): + sjoin_nearest(left_sgpd, right_sgpd, max_distance="1") + one_result = sjoin_nearest(left_sgpd, right_sgpd, exclusive=1) + numpy_bool_result = sjoin_nearest( + left_sgpd, + right_sgpd, + exclusive=np.bool_(True), + ) + self._assert_matches_geopandas( + one_result, + left.sjoin_nearest(right, exclusive=1), + ) + self._assert_matches_geopandas( + numpy_bool_result, + left.sjoin_nearest(right, exclusive=np.bool_(True)), + ) + with pytest.raises(ValueError, match="must be boolean"): + sjoin_nearest(left_sgpd, right_sgpd, exclusive=2) + with pytest.raises(TypeError, match="must be a string"): + sjoin_nearest(left_sgpd, right_sgpd, distance_col=1) + + geographic_left = self._to_sedona(left.to_crs(4326)) + geographic_right = self._to_sedona(right.to_crs(4326)) + with pytest.warns(UserWarning, match="geographic CRS") as warnings_seen: + sjoin_nearest(geographic_left, geographic_right) + assert ( + sum("geographic CRS" in str(warning.message) for warning in warnings_seen) + == 2 + ) + + mismatched = self._to_sedona(right.to_crs(32631)) + with pytest.warns(UserWarning, match="CRS mismatch"): + sjoin_nearest(left_sgpd, mismatched) + + no_crs = self._to_sedona(right.set_crs(None, allow_override=True)) + with pytest.warns(UserWarning, match="Right CRS: None"): + sjoin_nearest(left_sgpd, no_crs) diff --git a/python/tests/sql/test_dataframe_api.py b/python/tests/sql/test_dataframe_api.py index 8ca732b5b62..36e60272ce7 100644 --- a/python/tests/sql/test_dataframe_api.py +++ b/python/tests/sql/test_dataframe_api.py @@ -1652,6 +1652,12 @@ (stp.ST_EqualsExact, (None, "", 0.0)), (stp.ST_EqualsExact, ("", None, 0.0)), (stp.ST_EqualsExact, ("", "", None)), + (stp.ST_KNN, (None, "", 1)), + (stp.ST_KNN, ("", None, 1)), + (stp.ST_KNN, ("", "", 1.0)), + (stp.ST_KNN, ("", "", 1, 1.0)), + (stp.ST_KNN, ("", "", 1, False, 1.0)), + (stp.ST_KNN, ("", "", 1, False, True, 1.0)), (stp.ST_Overlaps, (None, "")), (stp.ST_Overlaps, ("", None)), (stp.ST_Touches, (None, "")), @@ -1962,6 +1968,75 @@ def run_spatial_query(): for future in futures: future.result() + def test_knn_query_local_options(self): + queries = ( + self.spark.createDataFrame([(0, "POINT (0 0)")], ["id", "wkt"]) + .select("id", stc.ST_GeomFromWKT("wkt").alias("geom")) + .alias("queries") + ) + tied_objects = ( + self.spark.createDataFrame( + [(0, "POINT (-1 0)"), (1, "POINT (1 0)"), (2, "POINT (10 0)")], + ["id", "wkt"], + ) + .select("id", stc.ST_GeomFromWKT("wkt").alias("geom")) + .alias("objects") + ) + + tied = queries.join( + tied_objects, + stp.ST_KNN( + queries["geom"], + tied_objects["geom"], + 1, + use_spheroid=False, + include_ties=True, + ), + ).select(tied_objects["id"].alias("id")) + assert sorted(row.id for row in tied.collect()) == [0, 1] + + for broadcast_side in ("queries", "objects"): + query_plan = ( + queries.hint("broadcast") if broadcast_side == "queries" else queries + ).alias("queries") + object_plan = ( + tied_objects.hint("broadcast") + if broadcast_side == "objects" + else tied_objects + ).alias("objects") + no_ties = query_plan.join( + object_plan, + stp.ST_KNN( + query_plan["geom"], + object_plan["geom"], + 1, + use_spheroid=False, + include_ties=False, + ), + ).select(object_plan["id"].alias("id")) + assert [row.id for row in no_ties.collect()] == [0] + + exclusive_objects = ( + self.spark.createDataFrame( + [(0, "POINT (0 0)"), (1, "POINT (2 0)")], + ["id", "wkt"], + ) + .select("id", stc.ST_GeomFromWKT("wkt").alias("geom")) + .alias("exclusive_objects") + ) + exclusive = queries.join( + exclusive_objects, + stp.ST_KNN( + queries["geom"], + exclusive_objects["geom"], + 1, + use_spheroid=False, + include_ties=True, + exclusive=True, + ), + ).select(exclusive_objects["id"].alias("id")) + assert [row.id for row in exclusive.collect()] == [1] + @pytest.mark.skipif( os.getenv("SPARK_REMOTE") is not None and pyspark.__version__ < "4.0.0", reason="DBSCAN requires checkpoint directory which is not available in Spark Connect mode before Spark 4.0.0", diff --git a/spark/common/src/main/java/org/apache/sedona/core/joinJudgement/InMemoryKNNJoinIterator.java b/spark/common/src/main/java/org/apache/sedona/core/joinJudgement/InMemoryKNNJoinIterator.java index 54ba42485ed..b913929345d 100644 --- a/spark/common/src/main/java/org/apache/sedona/core/joinJudgement/InMemoryKNNJoinIterator.java +++ b/spark/common/src/main/java/org/apache/sedona/core/joinJudgement/InMemoryKNNJoinIterator.java @@ -20,12 +20,16 @@ import java.util.ArrayList; import java.util.Collections; +import java.util.Comparator; +import java.util.IdentityHashMap; import java.util.Iterator; import java.util.LinkedHashSet; import java.util.List; import java.util.NoSuchElementException; +import java.util.Set; import org.apache.commons.lang3.tuple.Pair; import org.apache.sedona.core.enums.DistanceMetric; +import org.apache.sedona.core.wrapper.KnnGeometryMetadata; import org.apache.sedona.core.wrapper.UniqueGeometry; import org.apache.spark.util.LongAccumulator; import org.locationtech.jts.geom.Envelope; @@ -41,6 +45,8 @@ public class InMemoryKNNJoinIterator private final int k; private final DistanceMetric distanceMetric; private final boolean includeTies; + private final boolean exclusive; + private final boolean queryLocalTieSemantics; private final ItemDistance itemDistance; private final LongAccumulator streamCount; @@ -57,13 +63,57 @@ public InMemoryKNNJoinIterator( boolean includeTies, LongAccumulator streamCount, LongAccumulator resultCount) { + this( + querySideIterator, + strTree, + k, + distanceMetric, + includeTies, + false, + streamCount, + resultCount); + } + + public InMemoryKNNJoinIterator( + Iterator querySideIterator, + STRtree strTree, + int k, + DistanceMetric distanceMetric, + boolean includeTies, + boolean exclusive, + LongAccumulator streamCount, + LongAccumulator resultCount) { + this( + querySideIterator, + strTree, + k, + distanceMetric, + includeTies, + exclusive, + false, + streamCount, + resultCount); + } + + public InMemoryKNNJoinIterator( + Iterator querySideIterator, + STRtree strTree, + int k, + DistanceMetric distanceMetric, + boolean includeTies, + boolean exclusive, + boolean queryLocalTieSemantics, + LongAccumulator streamCount, + LongAccumulator resultCount) { this.querySideIterator = querySideIterator; this.strTree = strTree; this.k = k; this.distanceMetric = distanceMetric; this.includeTies = includeTies; - this.itemDistance = KnnJoinIndexJudgement.getItemDistance(distanceMetric); + this.exclusive = exclusive; + this.queryLocalTieSemantics = queryLocalTieSemantics; + this.itemDistance = KnnJoinIndexJudgement.getItemDistance(distanceMetric, exclusive); this.streamCount = streamCount; this.resultCount = resultCount; @@ -108,28 +158,46 @@ private void populateNextBatch() { Object[] localK = strTree.nearestNeighbour(queryGeom.getEnvelopeInternal(), queryGeom, itemDistance, k); - if (includeTies) { - localK = getUpdatedLocalKWithTies(queryGeom, localK, strTree); - } - + List nearest = new ArrayList<>(); for (Object obj : localK) { U candidate = (U) obj; + if (!exclusive || !KnnJoinIndexJudgement.areTopologicallyEqual(queryGeom, candidate)) { + nearest.add(candidate); + } + } + if (queryLocalTieSemantics) { + // Inspect every candidate tied at the local kth distance even when ties are excluded. Stable + // row identity can then select the same k candidates in regular and broadcast plans. + nearest = getUpdatedLocalKWithTies(queryGeom, nearest, strTree); + if (!includeTies) { + nearest = deterministicTopK(queryGeom, nearest); + } + } else if (includeTies) { + nearest = getUpdatedLocalKWithTies(queryGeom, nearest, strTree); + } + + for (U candidate : nearest) { Pair pair = Pair.of(queryItem, candidate); currentResults.add(pair); resultCount.add(1); } } - private Object[] getUpdatedLocalKWithTies( - Geometry streamShape, Object[] localK, STRtree strTree) { + private List getUpdatedLocalKWithTies(Geometry streamShape, List localK, STRtree strTree) { + if (localK.isEmpty()) { + return localK; + } + Envelope searchEnvelope = streamShape.getEnvelopeInternal(); // get the maximum distance from the k nearest neighbors double maxDistance = 0.0; - LinkedHashSet uniqueCandidates = new LinkedHashSet<>(); - for (Object obj : localK) { - U candidate = (U) obj; + Set uniqueCandidates = + queryLocalTieSemantics + ? Collections.newSetFromMap(new IdentityHashMap<>()) + : new LinkedHashSet<>(); + for (U candidate : localK) { uniqueCandidates.add(candidate); - double distance = streamShape.distance(candidate); + double distance = tieDistance(streamShape, candidate); if (distance > maxDistance) { maxDistance = distance; } @@ -138,18 +206,57 @@ private Object[] getUpdatedLocalKWithTies( List candidates = strTree.query(searchEnvelope); if (!candidates.isEmpty()) { // update localK with all candidates that are within the maxDistance - List tiedResults = new ArrayList<>(); + List tiedResults = new ArrayList<>(); // add all localK - Collections.addAll(tiedResults, localK); + tiedResults.addAll(localK); for (U candidate : candidates) { - double distance = streamShape.distance(candidate); - if (distance == maxDistance && !uniqueCandidates.contains(candidate)) { - tiedResults.add(candidate); + if (exclusive && KnnJoinIndexJudgement.areTopologicallyEqual(streamShape, candidate)) { + continue; + } + double distance = tieDistance(streamShape, candidate); + if (distance == maxDistance) { + // Keep the legacy equality-based membership check unchanged. Query-local semantics use + // identity so duplicate input rows with equal geometries remain separate candidates. + boolean isNewCandidate; + if (queryLocalTieSemantics) { + isNewCandidate = uniqueCandidates.add(candidate); + } else { + isNewCandidate = !uniqueCandidates.contains(candidate); + } + if (isNewCandidate) { + tiedResults.add(candidate); + } } } - localK = tiedResults.toArray(); + localK = tiedResults; } return localK; } + + private List deterministicTopK(Geometry query, List candidates) { + if (candidates.size() <= k) { + return candidates; + } + candidates.sort( + Comparator.comparingDouble((U candidate) -> tieDistance(query, candidate)) + .thenComparingLong(this::stableCandidateId)); + return new ArrayList<>(candidates.subList(0, k)); + } + + private long stableCandidateId(U candidate) { + Object userData = candidate.getUserData(); + if (!(userData instanceof KnnGeometryMetadata)) { + throw new IllegalStateException( + "Query-local planar KNN candidates must carry stable row metadata"); + } + return ((KnnGeometryMetadata) userData).getUniqueId(); + } + + private double tieDistance(Geometry query, Geometry candidate) { + if (queryLocalTieSemantics) { + return KnnJoinIndexJudgement.distanceByMetric(query, candidate, distanceMetric); + } + return query.distance(candidate); + } } diff --git a/spark/common/src/main/java/org/apache/sedona/core/joinJudgement/KnnJoinIndexJudgement.java b/spark/common/src/main/java/org/apache/sedona/core/joinJudgement/KnnJoinIndexJudgement.java index 0dda5869861..819a4153c29 100644 --- a/spark/common/src/main/java/org/apache/sedona/core/joinJudgement/KnnJoinIndexJudgement.java +++ b/spark/common/src/main/java/org/apache/sedona/core/joinJudgement/KnnJoinIndexJudgement.java @@ -32,6 +32,7 @@ import org.apache.spark.util.LongAccumulator; import org.locationtech.jts.geom.Geometry; import org.locationtech.jts.index.strtree.GeometryItemDistance; +import org.locationtech.jts.index.strtree.ItemBoundable; import org.locationtech.jts.index.strtree.ItemDistance; import org.locationtech.jts.index.strtree.STRtree; @@ -48,6 +49,8 @@ public class KnnJoinIndexJudgement private final int k; private final DistanceMetric distanceMetric; private final boolean includeTies; + private final boolean exclusive; + private final boolean queryLocalTieSemantics; private final Broadcast> broadcastQueryObjects; private final Broadcast broadcastObjectsTreeIndex; @@ -73,10 +76,62 @@ public KnnJoinIndexJudgement( LongAccumulator streamCount, LongAccumulator resultCount, LongAccumulator candidateCount) { + this( + k, + distanceMetric, + includeTies, + false, + broadcastQueryObjects, + broadcastObjectsTreeIndex, + buildCount, + streamCount, + resultCount, + candidateCount); + } + + public KnnJoinIndexJudgement( + int k, + DistanceMetric distanceMetric, + boolean includeTies, + boolean exclusive, + Broadcast> broadcastQueryObjects, + Broadcast broadcastObjectsTreeIndex, + LongAccumulator buildCount, + LongAccumulator streamCount, + LongAccumulator resultCount, + LongAccumulator candidateCount) { + this( + k, + distanceMetric, + includeTies, + exclusive, + false, + broadcastQueryObjects, + broadcastObjectsTreeIndex, + buildCount, + streamCount, + resultCount, + candidateCount); + } + + public KnnJoinIndexJudgement( + int k, + DistanceMetric distanceMetric, + boolean includeTies, + boolean exclusive, + boolean queryLocalTieSemantics, + Broadcast> broadcastQueryObjects, + Broadcast broadcastObjectsTreeIndex, + LongAccumulator buildCount, + LongAccumulator streamCount, + LongAccumulator resultCount, + LongAccumulator candidateCount) { super(null, buildCount, streamCount, resultCount, candidateCount); this.k = k; this.distanceMetric = distanceMetric; this.includeTies = includeTies; + this.exclusive = exclusive; + this.queryLocalTieSemantics = queryLocalTieSemantics; this.broadcastQueryObjects = broadcastQueryObjects; this.broadcastObjectsTreeIndex = broadcastObjectsTreeIndex; } @@ -104,7 +159,15 @@ public Iterator> call(Iterator queryShapes, Iterator objectShap STRtree strTree = buildSTRtree(objectShapes); return new InMemoryKNNJoinIterator<>( - queryShapes, strTree, k, distanceMetric, includeTies, streamCount, resultCount); + queryShapes, + strTree, + k, + distanceMetric, + includeTies, + exclusive, + queryLocalTieSemantics, + streamCount, + resultCount); } /** @@ -127,7 +190,15 @@ public Iterator> callUsingBroadcastObjectIndex(Iterator queryShape // broadcasted, the STRtree built from the broadcasted object should be able to fit into memory. STRtree strTree = broadcastObjectsTreeIndex.getValue(); return new InMemoryKNNJoinIterator<>( - queryShapes, strTree, k, distanceMetric, includeTies, streamCount, resultCount); + queryShapes, + strTree, + k, + distanceMetric, + includeTies, + exclusive, + queryLocalTieSemantics, + streamCount, + resultCount); } /** @@ -148,7 +219,15 @@ public Iterator> callUsingBroadcastQueryList(Iterator objectShapes List queryItems = broadcastQueryObjects.getValue(); STRtree strTree = buildSTRtree(objectShapes); return new InMemoryKNNJoinIterator<>( - queryItems.iterator(), strTree, k, distanceMetric, includeTies, streamCount, resultCount); + queryItems.iterator(), + strTree, + k, + distanceMetric, + includeTies, + exclusive, + queryLocalTieSemantics, + streamCount, + resultCount); } private STRtree buildSTRtree(Iterator objectShapes) { @@ -227,8 +306,41 @@ public static double distance( } public static ItemDistance getItemDistance(DistanceMetric distanceMetric) { - ItemDistance itemDistance; - itemDistance = getItemDistanceByMetric(distanceMetric); - return itemDistance; + return getItemDistance(distanceMetric, false); + } + + public static ItemDistance getItemDistance(DistanceMetric distanceMetric, boolean exclusive) { + ItemDistance itemDistance = getItemDistanceByMetric(distanceMetric); + return exclusive ? new ExclusiveItemDistance(itemDistance) : itemDistance; + } + + public static boolean areTopologicallyEqual(Geometry left, Geometry right) { + if (left == right) { + return true; + } + // Topologically equal geometries must have equal envelopes. This avoids constructing a DE-9IM + // relate matrix for nearly every candidate considered by an exclusive KNN search. + if (!left.getEnvelopeInternal().equals(right.getEnvelopeInternal())) { + return false; + } + return left.equalsTopo(right); + } + + private static class ExclusiveItemDistance implements ItemDistance, Serializable { + private final ItemDistance delegate; + + private ExclusiveItemDistance(ItemDistance delegate) { + this.delegate = delegate; + } + + @Override + public double distance(ItemBoundable item1, ItemBoundable item2) { + Geometry left = (Geometry) item1.getItem(); + Geometry right = (Geometry) item2.getItem(); + if (areTopologicallyEqual(left, right)) { + return Double.MAX_VALUE; + } + return delegate.distance(item1, item2); + } } } diff --git a/spark/common/src/main/java/org/apache/sedona/core/spatialOperator/JoinQuery.java b/spark/common/src/main/java/org/apache/sedona/core/spatialOperator/JoinQuery.java index 7b55dd07633..6ac6e74a63f 100644 --- a/spark/common/src/main/java/org/apache/sedona/core/spatialOperator/JoinQuery.java +++ b/spark/common/src/main/java/org/apache/sedona/core/spatialOperator/JoinQuery.java @@ -18,6 +18,7 @@ */ package org.apache.sedona.core.spatialOperator; +import java.io.Serializable; import java.util.*; import org.apache.commons.lang3.tuple.Pair; import org.apache.log4j.LogManager; @@ -32,6 +33,7 @@ import org.apache.sedona.core.spatialPartitioning.SpatialPartitioner; import org.apache.sedona.core.spatialRDD.CircleRDD; import org.apache.sedona.core.spatialRDD.SpatialRDD; +import org.apache.sedona.core.wrapper.KnnGeometryMetadata; import org.apache.sedona.core.wrapper.UniqueGeometry; import org.apache.spark.SparkContext; import org.apache.spark.api.java.JavaPairRDD; @@ -773,6 +775,61 @@ public static JavaPairRDD knnJoin boolean includeTies, boolean broadcastJoin) throws Exception { + return knnJoin( + queryRDD, objectRDD, joinParams, includeTies, false, broadcastJoin, false, false); + } + + /** + * Joins each query geometry to its nearest object geometries. + * + * @param queryRDD geometries for which neighbours are searched + * @param objectRDD candidate neighbour geometries + * @param joinParams KNN join parameters + * @param includeTies whether to return every candidate tied at the kth distance + * @param exclusive whether candidates topologically equal to the query are excluded before + * ranking + * @param broadcastJoin whether either side uses the broadcast KNN path + */ + public static JavaPairRDD knnJoin( + SpatialRDD queryRDD, + SpatialRDD objectRDD, + JoinParams joinParams, + boolean includeTies, + boolean exclusive, + boolean broadcastJoin) + throws Exception { + return knnJoin( + queryRDD, objectRDD, joinParams, includeTies, exclusive, broadcastJoin, false, false); + } + + /** + * Joins each query geometry to its nearest object geometries. + * + * @param queryRDD geometries for which neighbours are searched + * @param objectRDD candidate neighbour geometries + * @param joinParams KNN join parameters + * @param includeTies whether to return every candidate tied at the kth distance + * @param exclusive whether candidates topologically equal to the query are excluded before + * ranking + * @param broadcastJoin whether either side uses the broadcast KNN path + * @param mergeReplicatedQueries whether regular-plan local candidates are reconciled by stable + * row identity + * @param queryLocalTieSemantics whether tie distances use the requested metric and duplicate + * input rows retain their independent identity; callers enabling this internal execution mode + * must attach {@link KnnGeometryMetadata} to both inputs before partitioning + */ + public static JavaPairRDD knnJoin( + SpatialRDD queryRDD, + SpatialRDD objectRDD, + JoinParams joinParams, + boolean includeTies, + boolean exclusive, + boolean broadcastJoin, + boolean mergeReplicatedQueries, + boolean queryLocalTieSemantics) + throws Exception { + validateKnnExecutionOptions( + joinParams, broadcastJoin, mergeReplicatedQueries, queryLocalTieSemantics); verifyCRSMatch(queryRDD, objectRDD); if (!broadcastJoin) verifyPartitioningNumberMatch(queryRDD, objectRDD); @@ -811,7 +868,7 @@ public static JavaPairRDD knnJoin // The reason for using objectRDD as the right side is that the partitions are built on the // right side. - final JavaRDD> joinResult; + JavaRDD> joinResult; if (broadcastObjectsTreeIndex == null && broadcastQueryObjects == null) { // no broadcast join final KnnJoinIndexJudgement judgement = @@ -819,6 +876,8 @@ public static JavaPairRDD knnJoin joinParams.k, joinParams.distanceMetric, includeTies, + exclusive, + queryLocalTieSemantics, null, null, buildCount, @@ -827,6 +886,11 @@ public static JavaPairRDD knnJoin candidateCount); joinResult = queryRDD.spatialPartitionedRDD.zipPartitions(objectRDD.spatialPartitionedRDD, judgement); + if (mergeReplicatedQueries) { + joinResult = + mergeReplicatedKnnResults( + joinResult, joinParams.k, joinParams.distanceMetric, includeTies, exclusive); + } } else if (broadcastObjectsTreeIndex != null) { // broadcast join with objectRDD as broadcast side final KnnJoinIndexJudgement judgement = @@ -834,6 +898,8 @@ public static JavaPairRDD knnJoin joinParams.k, joinParams.distanceMetric, includeTies, + exclusive, + queryLocalTieSemantics, null, broadcastObjectsTreeIndex, buildCount, @@ -849,19 +915,201 @@ public static JavaPairRDD knnJoin joinParams.k, joinParams.distanceMetric, includeTies, + exclusive, + queryLocalTieSemantics, broadcastQueryObjects, null, buildCount, streamCount, resultCount, candidateCount); - joinResult = querySideBroadcastKNNJoin(objectRDD, joinParams, judgement, includeTies); + joinResult = + querySideBroadcastKNNJoin( + objectRDD, joinParams, judgement, includeTies, exclusive, queryLocalTieSemantics); } return joinResult.mapToPair( (PairFunction, U, T>) pair -> new Tuple2<>(pair.getKey(), pair.getValue())); } + static void validateKnnExecutionOptions( + JoinParams joinParams, + boolean broadcastJoin, + boolean mergeReplicatedQueries, + boolean queryLocalTieSemantics) { + if (mergeReplicatedQueries && broadcastJoin) { + throw new IllegalArgumentException( + "Replicated KNN reconciliation is only valid for the regular join plan"); + } + if (mergeReplicatedQueries && !queryLocalTieSemantics) { + throw new IllegalArgumentException( + "Replicated KNN reconciliation requires query-local tie semantics and stable row metadata"); + } + if (queryLocalTieSemantics && joinParams.distanceMetric != DistanceMetric.EUCLIDEAN) { + throw new IllegalArgumentException( + "Query-local tie semantics are only supported for planar Euclidean KNN joins"); + } + } + + /** + * Reconciles partition-local candidates for query geometries routed to multiple base cells. + * + *

Each input row has an independent stable ID in {@link KnnGeometryMetadata}. Object IDs + * remove only partition replicas, so distinct duplicate input rows remain distinct candidates. + */ + private static + JavaRDD> mergeReplicatedKnnResults( + JavaRDD> localResults, + int k, + DistanceMetric distanceMetric, + boolean includeTies, + boolean exclusive) { + return localResults + .mapToPair( + pair -> + new Tuple2>( + getKnnRowId(pair.getKey(), "query"), Pair.of(pair.getKey(), pair.getValue()))) + .combineByKey( + pair -> + new KnnCandidateAccumulator(k, distanceMetric, includeTies, exclusive) + .add(pair), + KnnCandidateAccumulator::add, + KnnCandidateAccumulator::merge) + .values() + .flatMap(KnnCandidateAccumulator::iterator); + } + + private static long getKnnRowId(Geometry geometry, String role) { + Object userData = geometry.getUserData(); + if (!(userData instanceof KnnGeometryMetadata)) { + throw new IllegalStateException( + "Query-local KNN " + + role + + " geometries must carry KnnGeometryMetadata; check execution option wiring"); + } + return ((KnnGeometryMetadata) userData).getUniqueId(); + } + + /** + * Mergeable bounded state for one replicated query row. + * + *

The map is ordered first by distance and then by stable object-row ID. It contains at most + * {@code k} candidates, except that include-ties mode retains every candidate at the current kth + * distance because all of them are required output. + */ + private static final class KnnCandidateAccumulator + implements Serializable { + private final int k; + private final DistanceMetric distanceMetric; + private final boolean includeTies; + private final boolean exclusive; + private final NavigableMap> candidates = new TreeMap<>(); + private final Set candidateIds = new HashSet<>(); + private U query; + private long size; + + private KnnCandidateAccumulator( + int k, DistanceMetric distanceMetric, boolean includeTies, boolean exclusive) { + this.k = k; + this.distanceMetric = distanceMetric; + this.includeTies = includeTies; + this.exclusive = exclusive; + } + + private KnnCandidateAccumulator add(Pair pair) { + if (query == null) { + query = pair.getKey(); + } + T candidate = pair.getValue(); + if (exclusive && KnnJoinIndexJudgement.areTopologicallyEqual(query, candidate)) { + return this; + } + + long candidateId = getKnnRowId(candidate, "candidate"); + if (candidateIds.contains(candidateId)) { + return this; + } + double distance = KnnJoinIndexJudgement.distance(query, candidate, distanceMetric); + + if (size < k) { + insert(distance, candidateId, candidate); + } else if (includeTies) { + addWithTies(distance, candidateId, candidate); + } else { + addWithoutTies(distance, candidateId, candidate); + } + return this; + } + + private void addWithTies(double distance, long candidateId, T candidate) { + double kthDistance = candidates.lastKey(); + int comparison = Double.compare(distance, kthDistance); + if (comparison > 0) { + return; + } + if (comparison == 0) { + insert(distance, candidateId, candidate); + return; + } + + long betterCandidateCount = size - candidates.lastEntry().getValue().size(); + insert(distance, candidateId, candidate); + if (betterCandidateCount + 1 >= k) { + removeDistance(kthDistance); + } + } + + private void addWithoutTies(double distance, long candidateId, T candidate) { + Map.Entry> lastEntry = candidates.lastEntry(); + int distanceComparison = Double.compare(distance, lastEntry.getKey()); + if (distanceComparison > 0) { + return; + } + if (distanceComparison == 0 && candidateId >= lastEntry.getValue().lastKey()) { + return; + } + + insert(distance, candidateId, candidate); + Map.Entry> largestDistance = candidates.lastEntry(); + Map.Entry removed = largestDistance.getValue().pollLastEntry(); + candidateIds.remove(removed.getKey()); + size--; + if (largestDistance.getValue().isEmpty()) { + candidates.remove(largestDistance.getKey()); + } + } + + private void insert(double distance, long candidateId, T candidate) { + candidates.computeIfAbsent(distance, ignored -> new TreeMap<>()).put(candidateId, candidate); + candidateIds.add(candidateId); + size++; + } + + private void removeDistance(double distance) { + NavigableMap removed = candidates.remove(distance); + if (removed != null) { + candidateIds.removeAll(removed.keySet()); + size -= removed.size(); + } + } + + private KnnCandidateAccumulator merge(KnnCandidateAccumulator other) { + for (NavigableMap distanceCandidates : other.candidates.values()) { + for (T candidate : distanceCandidates.values()) { + add(Pair.of(other.query, candidate)); + } + } + return this; + } + + private Iterator> iterator() { + return candidates.values().stream() + .flatMap(distanceCandidates -> distanceCandidates.values().stream()) + .map(candidate -> Pair.of(query, candidate)) + .iterator(); + } + } + /** * Performs a KNN join where the query side is broadcasted. * @@ -886,7 +1134,9 @@ JavaRDD> querySideBroadcastKNNJoin( SpatialRDD objectRDD, JoinParams joinParams, KnnJoinIndexJudgement, T> judgement, - boolean includeTies) { + boolean includeTies, + boolean exclusive, + boolean queryLocalTieSemantics) { final JavaRDD> joinResult; JavaRDD, T>> joinResultMapped = objectRDD.rawSpatialRDD.mapPartitions(judgement::callUsingBroadcastQueryList); @@ -907,7 +1157,11 @@ JavaRDD> querySideBroadcastKNNJoin( List> sortedPairs = new ArrayList<>(); for (Pair, T> p : values) { Pair newPair = Pair.of(p.getKey().getOriginalGeometry(), p.getValue()); - sortedPairs.add(newPair); + if (!exclusive + || !KnnJoinIndexJudgement.areTopologicallyEqual( + newPair.getKey(), newPair.getValue())) { + sortedPairs.add(newPair); + } } // Sort pairs based on the distance function between the two geometries @@ -919,7 +1173,13 @@ JavaRDD> querySideBroadcastKNNJoin( double distance2 = KnnJoinIndexJudgement.distance( p2.getKey(), p2.getValue(), distanceMetric); - return Double.compare(distance1, distance2); // Sort ascending by distance + int distanceComparison = Double.compare(distance1, distance2); + if (distanceComparison != 0 || !queryLocalTieSemantics) { + return distanceComparison; + } + return Long.compare( + getKnnRowId(p1.getValue(), "candidate"), + getKnnRowId(p2.getValue(), "candidate")); }); if (includeTies) { diff --git a/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioner.java b/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioner.java index 64ad97c038f..4d6d847103b 100644 --- a/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioner.java +++ b/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioner.java @@ -69,7 +69,17 @@ public QuadTreeRTPartitioner(ExtendedQuadTree extendedQuadTree) { } public QuadTreeRTPartitioner nonOverlappedPartitioner() { - ExtendedQuadTree nonOverlappedTree = new ExtendedQuadTree<>(extendedQuadTree, true); + return nonOverlappedPartitioner(false); + } + + /** + * Returns a query-side partitioner over the original cells. + * + * @param useFullGeometry whether non-point geometries are routed to every cell they intersect + */ + public QuadTreeRTPartitioner nonOverlappedPartitioner(boolean useFullGeometry) { + ExtendedQuadTree nonOverlappedTree = + new ExtendedQuadTree<>(extendedQuadTree, true, useFullGeometry); return new QuadTreeRTPartitioner(nonOverlappedTree); } diff --git a/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioning.java b/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioning.java index 0bc94a4a6a3..de29a79acab 100644 --- a/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioning.java +++ b/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioning.java @@ -79,6 +79,62 @@ public STRtree getMbrSpatialIndex() { * @return */ public STRtree buildSTRTree(List samples, int k) { + return buildSTRTree(samples, k, false); + } + + /** + * Builds the sample R-tree used to derive expanded KNN partition boundaries. + * + * @param samples sampled object envelopes + * @param k number of neighbouring sample envelopes to cover + * @param deduplicateSamples whether identical envelopes count once + * @return the partition-boundary R-tree + */ + public STRtree buildSTRTree(List samples, int k, boolean deduplicateSamples) { + return buildSTRTree(samples, k, deduplicateSamples, deduplicateSamples, false); + } + + /** + * Builds sample-driven KNN partition bounds with an optional safe fallback when the collected + * sample cannot bound the requested number of neighbours. + * + * @param samples sampled object envelopes + * @param k number of neighbouring sample envelopes to cover + * @param deduplicateSamples whether identical envelopes count once + * @param fallbackOnInsufficientSamples whether insufficient samples use the global boundary + * @return the partition-boundary R-tree + */ + public STRtree buildSTRTree( + List samples, + int k, + boolean deduplicateSamples, + boolean fallbackOnInsufficientSamples) { + return buildSTRTree( + samples, + k, + deduplicateSamples, + fallbackOnInsufficientSamples, + fallbackOnInsufficientSamples); + } + + /** + * Builds KNN partition bounds and independently controls the exact circle predicate used by the + * accurate planar path. + * + * @param samples sampled object envelopes + * @param k number of neighbouring sample envelopes to cover + * @param deduplicateSamples whether identical envelopes count once + * @param fallbackOnInsufficientSamples whether insufficient samples use the global boundary + * @param useExactCircleIntersection whether to avoid the historical polygonal circle + * approximation + * @return the partition-boundary R-tree + */ + public STRtree buildSTRTree( + List samples, + int k, + boolean deduplicateSamples, + boolean fallbackOnInsufficientSamples, + boolean useExactCircleIntersection) { // The partitioned MBRs mbrs = new HashMap<>(); @@ -99,15 +155,32 @@ public STRtree buildSTRTree(List samples, int k) { // Insert samples into an STR tree for k-nearest neighbor search STRtree sampleTree = new STRtree(); - for (Envelope sample : samples) { + Collection distanceSamples = + deduplicateSamples ? new LinkedHashSet<>(samples) : samples; + for (Envelope sample : distanceSamples) { // convert sample to a point Point point = geometryFactory.createPoint( new Coordinate(sample.centre().getX(), sample.centre().getY())); + point.setUserData(sample); sampleTree.insert(sample, point); } - processPartitions(partitionMBRs, mbrs, k, sampleTree, geometryFactory); + if (fallbackOnInsufficientSamples && distanceSamples.size() < k) { + // The sampled objects cannot provide a safe distance bound for the requested number of + // distinct neighbours. Replicate every object partition to every query partition instead of + // risking pruning the first non-equal neighbour. + Envelope globalBoundary = new Envelope(); + for (QuadRectangle quadRect : partitionMBRs) { + globalBoundary.expandToInclude(quadRect.getEnvelope()); + } + for (QuadRectangle quadRect : partitionMBRs) { + mbrs.put(quadRect.partitionId, Collections.singletonList(new Envelope(globalBoundary))); + } + } else { + processPartitions( + partitionMBRs, mbrs, k, sampleTree, geometryFactory, useExactCircleIntersection); + } // Construct a spatial index for the MBRs this.mbrSpatialIndex = new STRtree(); @@ -127,9 +200,25 @@ public void processPartitions( int k, STRtree sampleTree, GeometryFactory geometryFactory) { + processPartitions(partitionMBRs, mbrs, k, sampleTree, geometryFactory, false); + } + private void processPartitions( + List partitionMBRs, + Map> mbrs, + int k, + STRtree sampleTree, + GeometryFactory geometryFactory, + boolean useExactCircleIntersection) { for (QuadRectangle quadRect : partitionMBRs) { - processPartition(partitionMBRs, quadRect, mbrs, k, sampleTree, geometryFactory); + processPartition( + partitionMBRs, + quadRect, + mbrs, + k, + sampleTree, + geometryFactory, + useExactCircleIntersection); } } @@ -139,7 +228,8 @@ private void processPartition( Map> mbrs, int k, STRtree sampleTree, - GeometryFactory geometryFactory) { + GeometryFactory geometryFactory, + boolean useExactCircleIntersection) { Envelope partitionMBR = quadRect.getEnvelope(); @@ -155,7 +245,7 @@ private void processPartition( // Calculate the maximum distance from the centroid to the k-nearest neighbors in the samples double maxDistance = getMaxDistanceFromSamples(k, sampleTree, centroid); List intersectingMBRs = - getMBRIntersectEnvelopes(ui, maxDistance, centroidX, centroidY); + getMBRIntersectEnvelopes(ui, maxDistance, centroidX, centroidY, useExactCircleIntersection); mbrs.put(quadRect.partitionId, intersectingMBRs); } @@ -227,15 +317,40 @@ private static double getMaxDistanceFromSamples(int k, STRtree sampleTree, Point sampleTree.nearestNeighbour( centroid.getEnvelopeInternal(), centroid, new EuclideanItemDistance(), k); - // 4 - Calculate the distance to the farthest neighbor + // Each selected sample represents an arbitrary geometry inside its envelope. The distance to + // the envelope's farthest corner is an upper bound on the distance to that geometry. Taking + // the largest such bound for any k sampled objects therefore safely upper-bounds the kth + // object distance, including for non-point objects. double maxDistance = 0; for (Object neighbor : kNearestNeighbors) { if (neighbor instanceof Geometry) { - Envelope neighborEnvelope = ((Geometry) neighbor).getEnvelopeInternal(); - Coordinate neighborCoord = - new Coordinate(neighborEnvelope.centre().getX(), neighborEnvelope.centre().getY()); - Point neighborPoint = geometryFactory.createPoint(neighborCoord); - double distance = centroid.distance(neighborPoint); + Object userData = ((Geometry) neighbor).getUserData(); + Envelope neighborEnvelope = + userData instanceof Envelope + ? (Envelope) userData + : ((Geometry) neighbor).getEnvelopeInternal(); + double distance = + Math.max( + centroid + .getCoordinate() + .distance( + new Coordinate(neighborEnvelope.getMinX(), neighborEnvelope.getMinY())), + Math.max( + centroid + .getCoordinate() + .distance( + new Coordinate(neighborEnvelope.getMinX(), neighborEnvelope.getMaxY())), + Math.max( + centroid + .getCoordinate() + .distance( + new Coordinate( + neighborEnvelope.getMaxX(), neighborEnvelope.getMinY())), + centroid + .getCoordinate() + .distance( + new Coordinate( + neighborEnvelope.getMaxX(), neighborEnvelope.getMaxY()))))); if (distance > maxDistance) { maxDistance = distance; } @@ -257,7 +372,11 @@ private static double getMaxDistanceFromSamples(int k, STRtree sampleTree, Point * @return */ private List getMBRIntersectEnvelopes( - double ui, double maxDistance, double centroidX, double centroidY) { + double ui, + double maxDistance, + double centroidX, + double centroidY, + boolean useExactCircleIntersection) { // 5 - Construct the circle with radius ui and center centroid // Calculate the radius of the circle double gamma_i = 2 * ui + maxDistance; @@ -268,20 +387,32 @@ private List getMBRIntersectEnvelopes( centroidX - gamma_i, centroidX + gamma_i, centroidY - gamma_i, centroidY + gamma_i); - Coordinate center = new Coordinate(centroidX, centroidY); - Geometry circle = geometryFactory.createPoint(center).buffer(gamma_i); - // 6 - Compute all the MBRs that intersect with the circle and add them to a hash map List candidateEnvelopes = strTree.query(circleEnvelope); - // Filter the candidate envelopes to find those that intersect with the circle + // Filter using exact point-to-envelope distance. Geometry.buffer approximates a circle with an + // inscribed polygon and can incorrectly exclude an envelope close to the true circular arc. List intersectingMBRs = new ArrayList<>(); for (Envelope candidateEnvelope : candidateEnvelopes) { - Geometry envelopeGeometry = geometryFactory.toGeometry(candidateEnvelope); - if (circle.intersects(envelopeGeometry)) { + if (intersectsCircle( + candidateEnvelope, centroidX, centroidY, gamma_i, useExactCircleIntersection)) { intersectingMBRs.add(candidateEnvelope); } } return intersectingMBRs; } + + static boolean intersectsCircle( + Envelope envelope, + double centerX, + double centerY, + double radius, + boolean useExactCircleIntersection) { + Coordinate center = new Coordinate(centerX, centerY); + if (useExactCircleIntersection) { + return envelope.distance(new Envelope(center)) <= radius; + } + Geometry circle = geometryFactory.createPoint(center).buffer(radius); + return circle.intersects(geometryFactory.toGeometry(envelope)); + } } diff --git a/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/quadtree/ExtendedQuadTree.java b/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/quadtree/ExtendedQuadTree.java index 9925e93f8af..53c5c523d26 100644 --- a/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/quadtree/ExtendedQuadTree.java +++ b/spark/common/src/main/java/org/apache/sedona/core/spatialPartitioning/quadtree/ExtendedQuadTree.java @@ -62,6 +62,8 @@ public class ExtendedQuadTree extends PartitioningUtils implements Serializab private STRtree spatialExpandedBoundaryIndex; private boolean useNonOverlapped = false; + private boolean useFullGeometry = false; + private boolean useAccurateKnnPartitioning = false; public HashMap> getExpandedBoundaries() { return expandedBoundaries; @@ -89,12 +91,26 @@ public ExtendedQuadTree(Envelope boundary, int numPartitions) { * @param useNonOverlapped */ public ExtendedQuadTree(ExtendedQuadTree extendedQuadTree, boolean useNonOverlapped) { + this(extendedQuadTree, useNonOverlapped, false); + } + + /** + * Constructor to initialize the partitions list with non-overlapped boundaries. + * + * @param extendedQuadTree source tree + * @param useNonOverlapped whether original non-overlapped boundaries are used + * @param useFullGeometry whether a geometry is routed to every base cell it intersects + */ + public ExtendedQuadTree( + ExtendedQuadTree extendedQuadTree, boolean useNonOverlapped, boolean useFullGeometry) { this.boundary = extendedQuadTree.boundary; this.numPartitions = extendedQuadTree.numPartitions; this.expandedBoundaries = extendedQuadTree.expandedBoundaries; this.spatialExpandedBoundaryIndex = extendedQuadTree.spatialExpandedBoundaryIndex; this.partitionTree = extendedQuadTree.partitionTree; this.useNonOverlapped = useNonOverlapped; + this.useFullGeometry = useFullGeometry; + this.useAccurateKnnPartitioning = extendedQuadTree.useAccurateKnnPartitioning; } /** @@ -136,6 +152,10 @@ public Iterator> placeObject(Geometry geometry) { if (useNonOverlapped) { Objects.requireNonNull(geometry, "spatialObject"); + if (useFullGeometry && !(geometry instanceof Point)) { + return placeFullGeometry(geometry); + } + // KNN join uses geometry's centroid to calculate the distance final Envelope envelope = geometry.getCentroid().getEnvelopeInternal(); @@ -165,7 +185,10 @@ public Iterator> placeObject(Geometry geometry) { Envelope objectEnvelope = geometry.getEnvelopeInternal(); // Query the spatial index for intersecting envelopes - List intersectingIds = spatialExpandedBoundaryIndex.query(objectEnvelope); + Collection intersectingIds = spatialExpandedBoundaryIndex.query(objectEnvelope); + if (useAccurateKnnPartitioning) { + intersectingIds = new LinkedHashSet<>(intersectingIds); + } for (Integer partitionId : intersectingIds) { result.add(new Tuple2<>(partitionId, geometry)); @@ -175,6 +198,52 @@ public Iterator> placeObject(Geometry geometry) { } } + /** + * Routes a non-point query to each original quadtree cell that the geometry intersects. + * + *

The local KNN candidates from these cells are merged globally by query row identity. Using + * exact geometry/cell intersection avoids replicating a sparse line to every cell in its + * envelope. If a topology predicate cannot be evaluated, envelope routing is used as a safe + * fallback. + */ + private Iterator> placeFullGeometry(Geometry geometry) { + final Set partitionIds = getFullGeometryPartitionIds(geometry); + final List> result = new ArrayList<>(partitionIds.size()); + for (Integer partitionId : partitionIds) { + result.add(new Tuple2<>(partitionId, geometry)); + } + return result.iterator(); + } + + private Set getFullGeometryPartitionIds(Geometry geometry) { + final List matchedPartitions = + this.partitionTree.findZones(new QuadRectangle(geometry.getEnvelopeInternal())); + final Set result = new LinkedHashSet<>(); + + try { + for (QuadRectangle rectangle : matchedPartitions) { + Geometry cell = geometry.getFactory().toGeometry(rectangle.getEnvelope()); + if (geometry.intersects(cell)) { + result.add(rectangle.partitionId); + } + } + } catch (RuntimeException e) { + result.clear(); + for (QuadRectangle rectangle : matchedPartitions) { + result.add(rectangle.partitionId); + } + } + + // A non-empty geometry inside the global boundary should intersect at least one cell. Retain + // envelope matches if numerical precision at a cell edge produces no exact match. + if (result.isEmpty()) { + for (QuadRectangle rectangle : matchedPartitions) { + result.add(rectangle.partitionId); + } + } + return result; + } + /** * Check the geometry against the partition zones to find the IDs of overlapping ranges. Only * returns the IDs of the overlapping partitions. Note that the geometry can be in multiple ranges @@ -185,10 +254,15 @@ public Iterator> placeObject(Geometry geometry) { */ @Override public Set getKeys(Geometry geometry) { - if (!useNonOverlapped) { + if (useAccurateKnnPartitioning && useNonOverlapped) { + Objects.requireNonNull(geometry, "spatialObject"); + if (useFullGeometry && !(geometry instanceof Point)) { + return getFullGeometryPartitionIds(geometry); + } + // knn join uses the centroid of the geometry return partitionTree.getKeys(geometry.getCentroid()); - } else { + } else if (useAccurateKnnPartitioning) { // use the expanded boundaries Set keys = new HashSet<>(); Envelope objectEnvelope = geometry.getEnvelopeInternal(); @@ -198,6 +272,13 @@ public Set getKeys(Geometry geometry) { keys.addAll(intersectingIds); + return keys; + } else if (!useNonOverlapped) { + // Preserve the historical routing exposed by getKeys for legacy KNN partitioners. + return partitionTree.getKeys(geometry.getCentroid()); + } else { + Set keys = new HashSet<>(); + keys.addAll(spatialExpandedBoundaryIndex.query(geometry.getEnvelopeInternal())); return keys; } } @@ -223,6 +304,47 @@ public List fetchLeafZones() { * tree. */ public void build(int neighborSampleNumber) throws Exception { + build(neighborSampleNumber, false); + } + + /** + * Builds the partitioning tree and optionally counts identical sample envelopes once when + * deriving KNN distance bounds. + * + * @param neighborSampleNumber the number of neighbour samples + * @param deduplicateKnnSamples whether identical sample envelopes count once + */ + public void build(int neighborSampleNumber, boolean deduplicateKnnSamples) throws Exception { + build(neighborSampleNumber, deduplicateKnnSamples, deduplicateKnnSamples, false); + } + + /** + * Builds the partitioner and controls sample-driven KNN distance bounds independently from sample + * deduplication. + * + * @param neighborSampleNumber the number of neighbour samples + * @param deduplicateKnnSamples whether identical sample envelopes count once + * @param useKnnSamples whether samples are used and insufficient samples trigger a safe fallback + */ + public void build(int neighborSampleNumber, boolean deduplicateKnnSamples, boolean useKnnSamples) + throws Exception { + build(neighborSampleNumber, deduplicateKnnSamples, useKnnSamples, useKnnSamples); + } + + /** + * Builds the partitioner and independently gates the new accurate planar placement behavior. + * + * @param neighborSampleNumber the number of neighbour samples + * @param deduplicateKnnSamples whether identical sample envelopes count once + * @param useKnnSamples whether samples are used and insufficient samples trigger a safe fallback + * @param useAccurateKnnPartitioning whether to use exact bounds and replica-safe placement + */ + public void build( + int neighborSampleNumber, + boolean deduplicateKnnSamples, + boolean useKnnSamples, + boolean useAccurateKnnPartitioning) + throws Exception { // Force the quad-tree to grow up to a certain level // So the actual num of partitions might be slightly different int minLevel = (int) Math.max(Math.log(numPartitions) / Math.log(4), 0); @@ -231,9 +353,15 @@ public void build(int neighborSampleNumber) throws Exception { partitionTree = quadTreeRTPartitioning.getPartitionTree(); // Create the expanded boundaries - quadTreeRTPartitioning.buildSTRTree(samples, neighborSampleNumber); + quadTreeRTPartitioning.buildSTRTree( + samples, + neighborSampleNumber, + deduplicateKnnSamples, + useKnnSamples, + useAccurateKnnPartitioning); expandedBoundaries = quadTreeRTPartitioning.getMbrs(); spatialExpandedBoundaryIndex = quadTreeRTPartitioning.getMbrSpatialIndex(); + this.useAccurateKnnPartitioning = useAccurateKnnPartitioning; // Make sure not to broadcast all the samples used to build the Quad // tree to all nodes which are doing partitioning diff --git a/spark/common/src/main/java/org/apache/sedona/core/spatialRDD/SpatialRDD.java b/spark/common/src/main/java/org/apache/sedona/core/spatialRDD/SpatialRDD.java index 5d77e147660..ecb53e57867 100644 --- a/spark/common/src/main/java/org/apache/sedona/core/spatialRDD/SpatialRDD.java +++ b/spark/common/src/main/java/org/apache/sedona/core/spatialRDD/SpatialRDD.java @@ -101,6 +101,15 @@ public class SpatialRDD implements Serializable { /** The neighbor sample number. */ private int neighborSampleNumber = -1; + /** Whether KNN distance-bound samples with identical envelopes should be counted once. */ + private boolean deduplicateKnnSamples = false; + + /** Whether sampled object envelopes should be used to derive KNN partition distance bounds. */ + private boolean useKnnSamples = false; + + /** Whether the new exact planar KNN partitioning and replica handling should be used. */ + private boolean useAccurateKnnPartitioning = false; + public int getSampleNumber() { return sampleNumber; } @@ -123,6 +132,37 @@ public void setNeighborSampleNumber(int neighborSampleNumber) { this.neighborSampleNumber = neighborSampleNumber; } + /** + * Controls whether equal sample envelopes are counted once when deriving KNN partition bounds. + * + *

Exclusive KNN joins need this so an arbitrary number of geometries equal to a query cannot + * consume every neighbour sample and hide the next non-equal candidate outside the expanded + * partition. + * + * @param deduplicateKnnSamples whether to deduplicate sample envelopes + */ + public void setDeduplicateKnnSamples(boolean deduplicateKnnSamples) { + this.deduplicateKnnSamples = deduplicateKnnSamples; + } + + /** + * Controls whether object samples are used to derive safe KNN partition distance bounds. + * + * @param useKnnSamples whether to use sampled object envelopes + */ + public void setUseKnnSamples(boolean useKnnSamples) { + this.useKnnSamples = useKnnSamples; + } + + /** + * Controls exact bounds and replica-safe placement for planar query-local KNN. + * + * @param useAccurateKnnPartitioning whether to enable accurate planar partitioning + */ + public void setUseAccurateKnnPartitioning(boolean useAccurateKnnPartitioning) { + this.useAccurateKnnPartitioning = useAccurateKnnPartitioning; + } + /** * CRS transform. * @@ -312,7 +352,19 @@ public Envelope call(T geometry) throws Exception { { ExtendedQuadTree tree = new ExtendedQuadTree<>(paddedBoundary, numPartitions); ExtendedQuadTree extendedQuadTree = (ExtendedQuadTree) tree; - extendedQuadTree.build(neighborSampleNumber); + // Preserve the historical and geography ST_KNN partition bounds. Planar query-local + // 5/6-argument KNN opts into sample-driven bounds so replicated non-point queries can be + // merged exactly. + if (useKnnSamples) { + for (Envelope sample : samples) { + extendedQuadTree.insert(sample); + } + } + extendedQuadTree.build( + neighborSampleNumber, + deduplicateKnnSamples, + useKnnSamples, + useAccurateKnnPartitioning); partitioner = new QuadTreeRTPartitioner(extendedQuadTree); break; } diff --git a/spark/common/src/main/java/org/apache/sedona/core/wrapper/KnnGeometryMetadata.java b/spark/common/src/main/java/org/apache/sedona/core/wrapper/KnnGeometryMetadata.java new file mode 100644 index 00000000000..6bc06042bac --- /dev/null +++ b/spark/common/src/main/java/org/apache/sedona/core/wrapper/KnnGeometryMetadata.java @@ -0,0 +1,56 @@ +/* + * 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. + */ +package org.apache.sedona.core.wrapper; + +import java.io.Serializable; + +/** Stable row identity used by query-local planar KNN execution. */ +public final class KnnGeometryMetadata implements Serializable { + private final long uniqueId; + private final Object originalUserData; + + public KnnGeometryMetadata(long uniqueId, Object originalUserData) { + this.uniqueId = uniqueId; + this.originalUserData = originalUserData; + } + + /** + * Adds stable identity without nesting metadata when an execution path is prepared more than + * once. + */ + public static KnnGeometryMetadata wrap(long uniqueId, Object userData) { + if (userData instanceof KnnGeometryMetadata) { + return (KnnGeometryMetadata) userData; + } + return new KnnGeometryMetadata(uniqueId, userData); + } + + public long getUniqueId() { + return uniqueId; + } + + /** Returns the row payload, tolerating metadata nested by an older or external caller. */ + public Object getOriginalUserData() { + Object userData = originalUserData; + while (userData instanceof KnnGeometryMetadata) { + userData = ((KnnGeometryMetadata) userData).originalUserData; + } + return userData; + } +} diff --git a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Predicates.scala b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Predicates.scala index 0e693f9351c..37e151e9ca4 100644 --- a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Predicates.scala +++ b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/Predicates.scala @@ -350,7 +350,9 @@ private[apache] case class ST_3DDWithin(inputExpressions: Seq[Expression]) private[apache] case class ST_KNN(inputExpressions: Seq[Expression]) extends InferredExpression( inferrableFunction3(Predicates.knn), - inferrableFunction4(Predicates.knn)) { + inferrableFunction4(Predicates.knn), + inferrableFunction5(Predicates.knn), + inferrableFunction6(Predicates.knn)) { protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]) = { copy(inputExpressions = newChildren) diff --git a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_predicates.scala b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_predicates.scala index 33bc6898c4e..ab37950447b 100644 --- a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_predicates.scala +++ b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/expressions/st_predicates.scala @@ -93,4 +93,34 @@ object st_predicates { wrapExpression[ST_KNN](a, b, distance, useSphere) def ST_KNN(a: String, b: String, distance: Double, useSphere: Boolean): Column = wrapExpression[ST_KNN](a, b, distance, useSphere) + def ST_KNN( + a: Column, + b: Column, + distance: Column, + useSphere: Column, + includeTies: Column): Column = + wrapExpression[ST_KNN](a, b, distance, useSphere, includeTies) + def ST_KNN( + a: String, + b: String, + distance: Double, + useSphere: Boolean, + includeTies: Boolean): Column = + wrapExpression[ST_KNN](a, b, distance, useSphere, includeTies) + def ST_KNN( + a: Column, + b: Column, + distance: Column, + useSphere: Column, + includeTies: Column, + exclusive: Column): Column = + wrapExpression[ST_KNN](a, b, distance, useSphere, includeTies, exclusive) + def ST_KNN( + a: String, + b: String, + distance: Double, + useSphere: Boolean, + includeTies: Boolean, + exclusive: Boolean): Column = + wrapExpression[ST_KNN](a, b, distance, useSphere, includeTies, exclusive) } diff --git a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/BroadcastObjectSideKNNJoinExec.scala b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/BroadcastObjectSideKNNJoinExec.scala index f4bdae40d5d..6ddabe3e0c4 100644 --- a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/BroadcastObjectSideKNNJoinExec.scala +++ b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/BroadcastObjectSideKNNJoinExec.scala @@ -43,7 +43,9 @@ case class BroadcastObjectSideKNNJoinExec( spatialPredicate: SpatialPredicate, isGeography: Boolean, condition: Expression, - extraCondition: Option[Expression] = None) + extraCondition: Option[Expression] = None, + includeTies: Option[Boolean] = None, + exclusive: Boolean = false) extends SedonaBinaryExecNode with TraitKNNJoinQueryExec with Logging { @@ -77,7 +79,7 @@ case class BroadcastObjectSideKNNJoinExec( * @return */ def leftToSpatialRDD(rdd: RDD[UnsafeRow], shapeExpression: Expression): SpatialRDD[Geometry] = { - toSpatialRDD(rdd, shapeExpression) + toKNNSpatialRDD(rdd, shapeExpression) } /** @@ -93,7 +95,7 @@ case class BroadcastObjectSideKNNJoinExec( def rightToSpatialRDD( rdd: RDD[UnsafeRow], shapeExpression: Expression): SpatialRDD[Geometry] = { - toSpatialRDD(rdd, shapeExpression) + toKNNSpatialRDD(rdd, shapeExpression) } /** @@ -121,7 +123,8 @@ case class BroadcastObjectSideKNNJoinExec( require(numPartitions > 0, "The number of partitions must be greater than 0.") val kValue: Int = this.k.eval().asInstanceOf[Int] require(kValue >= 1, "The number of neighbors (k) must be equal or greater than 1.") - objectsShapes.setNeighborSampleNumber(kValue) + objectsShapes.setNeighborSampleNumber(if (exclusive) kValue + 1 else kValue) + objectsShapes.setDeduplicateKnnSamples(exclusive) broadcastJoin = true } diff --git a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/BroadcastQuerySideKNNJoinExec.scala b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/BroadcastQuerySideKNNJoinExec.scala index 575ded91251..6434eccdc98 100644 --- a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/BroadcastQuerySideKNNJoinExec.scala +++ b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/BroadcastQuerySideKNNJoinExec.scala @@ -44,7 +44,9 @@ case class BroadcastQuerySideKNNJoinExec( spatialPredicate: SpatialPredicate, isGeography: Boolean, condition: Expression, - extraCondition: Option[Expression] = None) + extraCondition: Option[Expression] = None, + includeTies: Option[Boolean] = None, + exclusive: Boolean = false) extends SedonaBinaryExecNode with TraitKNNJoinQueryExec with Logging { @@ -83,7 +85,7 @@ case class BroadcastQuerySideKNNJoinExec( rdd: RDD[UnsafeRow], shapeExpression: Expression, projection: Option[Seq[Expression]] = None): SpatialRDD[Geometry] = { - toSpatialRDD(rdd, shapeExpression) + toKNNSpatialRDD(rdd, shapeExpression) } /** @@ -100,7 +102,7 @@ case class BroadcastQuerySideKNNJoinExec( rdd: RDD[UnsafeRow], shapeExpression: Expression, projection: Option[Seq[Expression]] = None): SpatialRDD[Geometry] = { - toSpatialRDD(rdd, shapeExpression) + toKNNSpatialRDD(rdd, shapeExpression) } /** @@ -128,7 +130,8 @@ case class BroadcastQuerySideKNNJoinExec( require(numPartitions > 0, "The number of partitions must be greater than 0.") val kValue: Int = this.k.eval().asInstanceOf[Int] require(kValue >= 1, "The number of neighbors (k) must be equal or greater than 1.") - objectsShapes.setNeighborSampleNumber(kValue) + objectsShapes.setNeighborSampleNumber(if (exclusive) kValue + 1 else kValue) + objectsShapes.setDeduplicateKnnSamples(exclusive) // index the objects on regular partitions (not spatial partitions) // this avoids the cost of spatial partitioning diff --git a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/JoinQueryDetector.scala b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/JoinQueryDetector.scala index 68aad366b8b..70e261a2ab9 100644 --- a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/JoinQueryDetector.scala +++ b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/JoinQueryDetector.scala @@ -627,7 +627,7 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { rightShape, spatialPredicate = SpatialPredicate.KNN, isGeography = false, - condition, + extraCondition, Some(k))) case ST_KNN(Seq(leftShape, rightShape, k, useSpheroid)) => @@ -640,7 +640,31 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { rightShape, spatialPredicate = SpatialPredicate.KNN, isGeography = useSpheroidUnwrapped, - condition, + extraCondition, + Some(k))) + case ST_KNN(Seq(leftShape, rightShape, k, useSpheroid, _)) => + val useSpheroidUnwrapped = useSpheroid.eval().asInstanceOf[Boolean] + Some( + JoinQueryDetection( + left, + right, + leftShape, + rightShape, + spatialPredicate = SpatialPredicate.KNN, + isGeography = useSpheroidUnwrapped, + extraCondition, + Some(k))) + case ST_KNN(Seq(leftShape, rightShape, k, useSpheroid, _, _)) => + val useSpheroidUnwrapped = useSpheroid.eval().asInstanceOf[Boolean] + Some( + JoinQueryDetection( + left, + right, + leftShape, + rightShape, + spatialPredicate = SpatialPredicate.KNN, + isGeography = useSpheroidUnwrapped, + extraCondition, Some(k))) case _ => None @@ -665,7 +689,8 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { detection.isGeography, detection.extraCondition, detection.distance, - detection.geographyShape) + detection.geographyShape, + condition) case _ => Nil } @@ -831,6 +856,7 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { logInfo( "Planning knn join, left side is for queries and right size is for the object to be searched") + val options = extractKNNOptions(condition) KNNJoinExec( planLater(left), planLater(right), @@ -841,7 +867,9 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { spatialPredicate = null, isGeography, condition, - extractExtraKNNJoinCondition(condition)) :: Nil + extraCondition, + options.includeTies, + options.exclusive) :: Nil } private def planDistanceJoin( @@ -903,22 +931,39 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { } } - private def extractExtraKNNJoinCondition(condition: Expression): Option[Expression] = { - condition match { - case and: And => - // Check both left and right sides for ST_KNN or ST_AKNN - if (and.left.isInstanceOf[ST_KNN]) { - Some(and.right) - } else if (and.right.isInstanceOf[ST_KNN]) { - Some(and.left) - } else { - None - } - case _: ST_KNN => - None - case _ => - Some(condition) + private case class KNNOptions(includeTies: Option[Boolean] = None, exclusive: Boolean = false) + + private def extractKNNOptions(condition: Expression): KNNOptions = { + val knn = findKNNExpression(condition).getOrElse( + throw new IllegalArgumentException("KNN join condition does not contain ST_KNN")) + knn.inputExpressions match { + case Seq(_, _, _) | Seq(_, _, _, _) => + KNNOptions() + case Seq(_, _, _, _, includeTies) => + KNNOptions(includeTies = Some(evalKNNBoolean(includeTies, "includeTies"))) + case Seq(_, _, _, _, includeTies, exclusive) => + KNNOptions( + includeTies = Some(evalKNNBoolean(includeTies, "includeTies")), + exclusive = evalKNNBoolean(exclusive, "exclusive")) + case expressions => + throw new IllegalArgumentException( + s"ST_KNN expects between 3 and 6 arguments, got ${expressions.size}") + } + } + + private def findKNNExpression(expression: Expression): Option[ST_KNN] = + expression match { + case knn: ST_KNN => Some(knn) + case And(left, right) => + findKNNExpression(left).orElse(findKNNExpression(right)) + case _ => None } + + private def evalKNNBoolean(expression: Expression, name: String): Boolean = { + require(expression.foldable, s"ST_KNN $name must be a constant boolean") + val value = expression.eval() + require(value != null, s"ST_KNN $name must not be null") + value.asInstanceOf[Boolean] } private def planBroadcastJoin( @@ -933,7 +978,8 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { isGeography: Boolean, extraCondition: Option[Expression], distance: Option[Expression], - geographyShape: Boolean = false): Seq[SparkPlan] = { + geographyShape: Boolean = false, + condition: Option[Expression] = None): Seq[SparkPlan] = { val broadcastSide = joinType match { case Inner if broadcastLeft => Some(LeftSide) @@ -951,6 +997,8 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { if (spatialPredicate == SpatialPredicate.KNN) { { + val knnCondition = condition.get + val options = extractKNNOptions(knnCondition) // validate the k value for KNN join val kValue: Int = distance.get.eval().asInstanceOf[Int] require(kValue >= 1, "The number of neighbors (k) must be equal or greater than 1.") @@ -980,7 +1028,9 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { spatialPredicate, isGeography, condition = null, - extraCondition = None) :: Nil + extraCondition = extraCondition, + includeTies = options.includeTies, + exclusive = options.exclusive) :: Nil } else { // broadcast is on object side return BroadcastObjectSideKNNJoinExec( @@ -995,7 +1045,9 @@ class JoinQueryDetector(sparkSession: SparkSession) extends SparkStrategy { spatialPredicate, isGeography, condition = null, - extraCondition = None) :: Nil + extraCondition = extraCondition, + includeTies = options.includeTies, + exclusive = options.exclusive) :: Nil } } } diff --git a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/KNNJoinExec.scala b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/KNNJoinExec.scala index a879447d409..133cd886c3d 100644 --- a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/KNNJoinExec.scala +++ b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/KNNJoinExec.scala @@ -70,7 +70,9 @@ case class KNNJoinExec( spatialPredicate: SpatialPredicate, isGeography: Boolean, condition: Expression, - extraCondition: Option[Expression] = None) + extraCondition: Option[Expression] = None, + includeTies: Option[Boolean] = None, + exclusive: Boolean = false) extends SedonaBinaryExecNode with TraitKNNJoinQueryExec with Logging { @@ -109,7 +111,7 @@ case class KNNJoinExec( rdd: RDD[UnsafeRow], shapeExpression: Expression, projection: Option[Seq[Expression]] = None): SpatialRDD[Geometry] = { - toSpatialRDD(rdd, shapeExpression) + toKNNSpatialRDD(rdd, shapeExpression) } /** @@ -126,7 +128,7 @@ case class KNNJoinExec( rdd: RDD[UnsafeRow], shapeExpression: Expression, projection: Option[Seq[Expression]] = None): SpatialRDD[Geometry] = { - toSpatialRDD(rdd, shapeExpression) + toKNNSpatialRDD(rdd, shapeExpression) } /** @@ -163,11 +165,20 @@ case class KNNJoinExec( require(numPartitions > 0, "The number of partitions must be greater than 0.") val kValue: Int = this.k.eval().asInstanceOf[Int] require(kValue >= 1, "The number of neighbors (k) must be equal or greater than 1.") - objectsShapes.setNeighborSampleNumber(kValue) + objectsShapes.setNeighborSampleNumber(if (exclusive) kValue + 1 else kValue) + objectsShapes.setUseKnnSamples(useAccurateRegularKnn) + objectsShapes.setUseAccurateKnnPartitioning(useAccurateRegularKnn) + objectsShapes.setDeduplicateKnnSamples(exclusive) exactSpatialPartitioning(objectsShapes, queryShapes, numPartitions) } + /** + * The planar five- and six-argument forms opt into globally reconciled regular-plan results. + * Legacy and geography forms retain their historical centroid routing. + */ + override protected def useAccurateRegularKnn: Boolean = useQueryLocalTieSemantics + /** * Exact spatial partitioning for KNN join * @param dominantShapes @@ -179,10 +190,6 @@ case class KNNJoinExec( dominantShapes: SpatialRDD[Geometry], followerShapes: SpatialRDD[Geometry], numPartitions: Integer): Unit = { - // analyze the both RDDs to get the statistics (e.g., boundary) - dominantShapes.analyze() - followerShapes.analyze() - // expand the boundary for partition to include both RDDs dominantShapes.boundaryEnvelope.expandToInclude(followerShapes.boundaryEnvelope) @@ -191,7 +198,7 @@ case class KNNJoinExec( followerShapes.spatialPartitioning( dominantShapes.getPartitioner .asInstanceOf[QuadTreeRTPartitioner] - .nonOverlappedPartitioner()) + .nonOverlappedPartitioner(useAccurateRegularKnn)) dominantShapes.buildIndex(IndexType.RTREE, true) } diff --git a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/TraitKNNJoinQueryExec.scala b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/TraitKNNJoinQueryExec.scala index 5eba2a94e11..9cab478b916 100644 --- a/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/TraitKNNJoinQueryExec.scala +++ b/spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/strategy/join/TraitKNNJoinQueryExec.scala @@ -22,7 +22,9 @@ import org.apache.commons.lang3.Range import org.apache.sedona.core.spatialOperator.JoinQuery import org.apache.sedona.core.spatialOperator.JoinQuery.JoinParams import org.apache.sedona.core.spatialPartitioning.{QuadTreeRTPartitioner, SpatialPartitioner} +import org.apache.sedona.core.spatialRDD.SpatialRDD import org.apache.sedona.core.utils.SedonaConf +import org.apache.sedona.core.wrapper.KnnGeometryMetadata import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.codegen.GenerateUnsafeRowJoiner @@ -48,6 +50,12 @@ trait TraitKNNJoinQueryExec extends TraitJoinQueryExec { protected var broadcastJoin: Boolean = false protected var querySide: JoinSide = null + def includeTies: Option[Boolean] + def exclusive: Boolean + def isGeography: Boolean + protected def useAccurateRegularKnn: Boolean = false + protected def useQueryLocalTieSemantics: Boolean = + includeTies.isDefined && !isGeography private lazy val sedonaConf = SedonaConf.fromActiveSession override lazy val metrics: Map[String, SQLMetric] = Map.empty @@ -91,7 +99,16 @@ trait TraitKNNJoinQueryExec extends TraitJoinQueryExec { val (queryShapes, objectShapes) = toSpatialRddPair(queryResultsRaw, boundQueryShape, objectResultsRaw, boundObjectShape) + if (useQueryLocalTieSemantics) { + ensureStableRowIds(queryShapes) + ensureStableRowIds(objectShapes) + } + objectShapes.analyze() + queryShapes.analyze() + if (objectShapes.approximateTotalCount == 0 || queryShapes.approximateTotalCount == 0) { + return joinedRddToRowRdd(sparkContext.emptyRDD[(Geometry, Geometry)], swapped) + } log.info( "[SedonaSQL] Number of partitions on the objectShapes (right): " + objectResultsRaw.partitions.size) @@ -132,8 +149,11 @@ trait TraitKNNJoinQueryExec extends TraitJoinQueryExec { queryShapes, objectShapes, joinParams, - sedonaConf.isIncludeTieBreakersInKNNJoins, - broadcastJoin) + includeTies.getOrElse(sedonaConf.isIncludeTieBreakersInKNNJoins), + exclusive, + broadcastJoin, + useAccurateRegularKnn, + useQueryLocalTieSemantics) .rdd } else { sparkContext.parallelize(Seq[(Geometry, Geometry)]()) @@ -144,8 +164,11 @@ trait TraitKNNJoinQueryExec extends TraitJoinQueryExec { queryShapes, objectShapes, joinParams, - sedonaConf.isIncludeTieBreakersInKNNJoins, - broadcastJoin) + includeTies.getOrElse(sedonaConf.isIncludeTieBreakersInKNNJoins), + exclusive, + broadcastJoin, + useAccurateRegularKnn, + useQueryLocalTieSemantics) .rdd } @@ -153,6 +176,41 @@ trait TraitKNNJoinQueryExec extends TraitJoinQueryExec { joinedRddToRowRdd(matchesRDD, swapped) } + /** + * Converts rows to KNN shapes while dropping null and empty geometries. Such rows cannot + * participate in an inner KNN match, and removing them lets the executor detect an empty + * query/candidate side before attempting spatial partitioning. + */ + protected def toKNNSpatialRDD( + rdd: RDD[UnsafeRow], + shapeExpression: Expression): SpatialRDD[Geometry] = { + val spatialRdd = new SpatialRDD[Geometry] + spatialRdd.setRawSpatialRDD( + rdd + .flatMap { row => + val shape = TraitJoinQueryBase.shapeToGeometry(shapeExpression, row) + if (shape == null || shape.isEmpty) { + None + } else { + shape.setUserData(row.copy) + Some(shape) + } + } + .toJavaRDD()) + spatialRdd + } + + private def ensureStableRowIds(spatialRdd: SpatialRDD[Geometry]): Unit = { + spatialRdd.setRawSpatialRDD( + spatialRdd.rawSpatialRDD.rdd + .zipWithUniqueId() + .map { case (geometry, uniqueId) => + geometry.setUserData(KnnGeometryMetadata.wrap(uniqueId, geometry.getUserData)) + geometry + } + .toJavaRDD()) + } + def knnJoinPartitionNumOptimizer( objectSidePartNum: Int, querySidePartNum: Int, @@ -228,8 +286,8 @@ trait TraitKNNJoinQueryExec extends TraitJoinQueryExec { } val joined = iter.map { case (l, r) => - val leftRow = l.getUserData.asInstanceOf[UnsafeRow] - val rightRow = r.getUserData.asInstanceOf[UnsafeRow] + val leftRow = getOriginalRow(l) + val rightRow = getOriginalRow(r) if (swapped) joinRow(rightRow, leftRow) else @@ -246,6 +304,25 @@ trait TraitKNNJoinQueryExec extends TraitJoinQueryExec { } } + private def getOriginalRow(geometry: Geometry): UnsafeRow = { + geometry.getUserData match { + case metadata: KnnGeometryMetadata => + metadata.getOriginalUserData match { + case row: UnsafeRow => row + case other => + throw new IllegalStateException( + s"KNN stable row metadata must wrap UnsafeRow, found ${userDataType(other)}") + } + case row: UnsafeRow => row + case other => + throw new IllegalStateException( + s"KNN geometry must carry UnsafeRow user data, found ${userDataType(other)}") + } + } + + private def userDataType(userData: Any): String = + if (userData == null) "null" else userData.getClass.getName + private def saveKNNPartitionerToFile( partitioner: SpatialPartitioner, savePath: String): Unit = { diff --git a/spark/common/src/test/java/org/apache/sedona/core/joinJudgement/InMemoryKNNJoinIteratorTest.java b/spark/common/src/test/java/org/apache/sedona/core/joinJudgement/InMemoryKNNJoinIteratorTest.java new file mode 100644 index 00000000000..f48bab23fe6 --- /dev/null +++ b/spark/common/src/test/java/org/apache/sedona/core/joinJudgement/InMemoryKNNJoinIteratorTest.java @@ -0,0 +1,184 @@ +/* + * 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. + */ +package org.apache.sedona.core.joinJudgement; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotSame; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertTrue; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.sedona.core.enums.DistanceMetric; +import org.apache.sedona.core.wrapper.KnnGeometryMetadata; +import org.apache.spark.util.LongAccumulator; +import org.junit.Test; +import org.locationtech.jts.geom.Coordinate; +import org.locationtech.jts.geom.GeometryFactory; +import org.locationtech.jts.geom.LineString; +import org.locationtech.jts.geom.Point; +import org.locationtech.jts.index.strtree.STRtree; + +public class InMemoryKNNJoinIteratorTest { + private final GeometryFactory factory = new GeometryFactory(); + + @Test + public void queryLocalTieSemanticsPreserveDuplicateInputRows() { + Point query = point(0, 0); + Point duplicateA = point(1, 0); + Point duplicateB = point(1, 0); + + List results = + collect(query, Arrays.asList(duplicateA, duplicateB), DistanceMetric.EUCLIDEAN, true); + + assertEquals(2, results.size()); + assertNotSame(results.get(0), results.get(1)); + } + + @Test + public void legacyConstructorRetainsGeometryEqualityDeduplication() { + Point query = point(0, 0); + Point duplicateA = point(1, 0); + Point duplicateB = point(1, 0); + STRtree index = index(Arrays.asList(duplicateA, duplicateB)); + + InMemoryKNNJoinIterator iterator = + new InMemoryKNNJoinIterator<>( + Collections.singletonList(query).iterator(), + index, + 1, + DistanceMetric.EUCLIDEAN, + true, + new LongAccumulator(), + new LongAccumulator()); + + assertEquals(1, collectValues(iterator).size()); + } + + @Test + public void geographyRetainsHistoricalPlanarTieExpansion() { + Point query = point(0, 90); + List antipodes = Arrays.asList(point(0, -90), point(100, -90), point(100, -90)); + + List results = collect(query, antipodes, DistanceMetric.HAVERSINE, false); + + assertEquals(1, results.size()); + } + + @Test + public void queryLocalNoTiesSelectsLowestStableRowId() { + Point query = point(0, 0); + Point higherId = point(-1, 0); + higherId.setUserData(KnnGeometryMetadata.wrap(9, "higher")); + Point lowerId = point(1, 0); + lowerId.setUserData(KnnGeometryMetadata.wrap(2, "lower")); + + List results = + collect(query, Arrays.asList(higherId, lowerId), DistanceMetric.EUCLIDEAN, false, true); + + assertEquals(Collections.singletonList(lowerId), results); + } + + @Test + public void queryLocalNoTiesDoesNotRequireMetadataWithoutTieAmbiguity() { + Point query = point(0, 0); + Point candidate = point(1, 0); + + List results = + collect(query, Collections.singletonList(candidate), DistanceMetric.EUCLIDEAN, false, true); + + assertEquals(Collections.singletonList(candidate), results); + } + + @Test + public void stableMetadataWrappingIsIdempotentAndRobustToOlderNesting() { + Object originalRow = new Object(); + KnnGeometryMetadata metadata = KnnGeometryMetadata.wrap(4, originalRow); + + assertSame(metadata, KnnGeometryMetadata.wrap(7, metadata)); + assertSame(originalRow, new KnnGeometryMetadata(8, metadata).getOriginalUserData()); + } + + @Test + public void topologicalEqualityUsesSafeShortCircuits() { + Point point = point(0, 0); + assertTrue(KnnJoinIndexJudgement.areTopologicallyEqual(point, point)); + assertFalse(KnnJoinIndexJudgement.areTopologicallyEqual(point, point(1, 0))); + + LineString forward = + factory.createLineString(new Coordinate[] {new Coordinate(0, 0), new Coordinate(1, 0)}); + LineString reverse = + factory.createLineString(new Coordinate[] {new Coordinate(1, 0), new Coordinate(0, 0)}); + assertTrue(KnnJoinIndexJudgement.areTopologicallyEqual(forward, reverse)); + } + + private List collect( + Point query, + List candidates, + DistanceMetric distanceMetric, + boolean queryLocalTieSemantics) { + return collect(query, candidates, distanceMetric, true, queryLocalTieSemantics); + } + + private List collect( + Point query, + List candidates, + DistanceMetric distanceMetric, + boolean includeTies, + boolean queryLocalTieSemantics) { + InMemoryKNNJoinIterator iterator = + new InMemoryKNNJoinIterator<>( + Collections.singletonList(query).iterator(), + index(candidates), + 1, + distanceMetric, + includeTies, + false, + queryLocalTieSemantics, + new LongAccumulator(), + new LongAccumulator()); + return collectValues(iterator); + } + + private List collectValues(InMemoryKNNJoinIterator iterator) { + List results = new ArrayList<>(); + while (iterator.hasNext()) { + Pair pair = iterator.next(); + results.add(pair.getValue()); + } + return results; + } + + private STRtree index(List points) { + STRtree index = new STRtree(); + for (Point point : points) { + index.insert(point.getEnvelopeInternal(), point); + } + index.build(); + return index; + } + + private Point point(double x, double y) { + return factory.createPoint(new Coordinate(x, y)); + } +} diff --git a/spark/common/src/test/java/org/apache/sedona/core/spatialOperator/KnnExecutionOptionsTest.java b/spark/common/src/test/java/org/apache/sedona/core/spatialOperator/KnnExecutionOptionsTest.java new file mode 100644 index 00000000000..84ab5598f3f --- /dev/null +++ b/spark/common/src/test/java/org/apache/sedona/core/spatialOperator/KnnExecutionOptionsTest.java @@ -0,0 +1,106 @@ +/* + * 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. + */ +package org.apache.sedona.core.spatialOperator; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import java.util.Arrays; +import java.util.Collections; +import org.apache.sedona.core.TestBase; +import org.apache.sedona.core.enums.DistanceMetric; +import org.apache.sedona.core.enums.IndexType; +import org.apache.sedona.core.spatialOperator.JoinQuery.JoinParams; +import org.apache.sedona.core.spatialRDD.SpatialRDD; +import org.apache.spark.api.java.JavaPairRDD; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; +import org.locationtech.jts.geom.Coordinate; +import org.locationtech.jts.geom.GeometryFactory; +import org.locationtech.jts.geom.Point; + +public class KnnExecutionOptionsTest extends TestBase { + @BeforeClass + public static void setup() { + initialize(KnnExecutionOptionsTest.class.getName()); + } + + @AfterClass + public static void teardown() { + sc.stop(); + } + + @Test + public void acceptsSupportedPlanAndMetricCombinations() { + JoinQuery.validateKnnExecutionOptions(planarParams(), false, true, true); + JoinQuery.validateKnnExecutionOptions(planarParams(), true, false, true); + + JoinParams geographyParams = + new JoinParams(true, null, IndexType.RTREE, null, 1, DistanceMetric.HAVERSINE); + JoinQuery.validateKnnExecutionOptions(geographyParams, false, false, false); + JoinQuery.validateKnnExecutionOptions(geographyParams, true, false, false); + } + + @Test + public void replicatedReconciliationRejectsBroadcastPlan() { + assertThrows( + IllegalArgumentException.class, + () -> JoinQuery.validateKnnExecutionOptions(planarParams(), true, true, true)); + } + + @Test + public void replicatedReconciliationRequiresQueryLocalSemantics() { + assertThrows( + IllegalArgumentException.class, + () -> JoinQuery.validateKnnExecutionOptions(planarParams(), false, true, false)); + } + + @Test + public void queryLocalSemanticsRejectNonPlanarMetric() { + JoinParams geographyParams = + new JoinParams(true, null, IndexType.RTREE, null, 1, DistanceMetric.HAVERSINE); + assertThrows( + IllegalArgumentException.class, + () -> JoinQuery.validateKnnExecutionOptions(geographyParams, false, false, true)); + } + + @Test + public void convenienceOverloadDoesNotRequireInternalRowMetadata() throws Exception { + GeometryFactory factory = new GeometryFactory(); + SpatialRDD queries = new SpatialRDD<>(); + queries.setRawSpatialRDD( + sc.parallelize(Collections.singletonList(factory.createPoint(new Coordinate(0.0, 0.0))))); + SpatialRDD objects = new SpatialRDD<>(); + objects.setRawSpatialRDD( + sc.parallelize( + Arrays.asList( + factory.createPoint(new Coordinate(-1.0, 0.0)), + factory.createPoint(new Coordinate(1.0, 0.0))))); + + JavaPairRDD matches = + JoinQuery.knnJoin(queries, objects, planarParams(), false, false, true); + + assertEquals(1L, matches.count()); + } + + private JoinParams planarParams() { + return new JoinParams(true, null, IndexType.RTREE, null, 1, DistanceMetric.EUCLIDEAN); + } +} diff --git a/spark/common/src/test/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioningTest.java b/spark/common/src/test/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioningTest.java new file mode 100644 index 00000000000..7ecb0b0f903 --- /dev/null +++ b/spark/common/src/test/java/org/apache/sedona/core/spatialPartitioning/QuadTreeRTPartitioningTest.java @@ -0,0 +1,48 @@ +/* + * 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. + */ +package org.apache.sedona.core.spatialPartitioning; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; +import org.locationtech.jts.geom.Coordinate; +import org.locationtech.jts.geom.Envelope; +import org.locationtech.jts.geom.GeometryFactory; + +public class QuadTreeRTPartitioningTest { + + @Test + public void testCircleIntersectionDoesNotUsePolygonalBufferApproximation() { + double radius = 10.0; + double angle = Math.PI / 32.0; + Coordinate coordinate = new Coordinate(9.98 * Math.cos(angle), 9.98 * Math.sin(angle)); + Envelope candidate = new Envelope(coordinate); + + assertTrue(QuadTreeRTPartitioning.intersectsCircle(candidate, 0.0, 0.0, radius, true)); + assertFalse(QuadTreeRTPartitioning.intersectsCircle(candidate, 0.0, 0.0, radius, false)); + + GeometryFactory factory = new GeometryFactory(); + assertFalse( + factory + .createPoint(new Coordinate(0.0, 0.0)) + .buffer(radius) + .intersects(factory.createPoint(coordinate))); + } +} diff --git a/spark/common/src/test/java/org/apache/sedona/core/spatialPartitioning/quadtree/QuadTreePartitioningTest.java b/spark/common/src/test/java/org/apache/sedona/core/spatialPartitioning/quadtree/QuadTreePartitioningTest.java index 42dec5c34ef..f53a5de40ff 100644 --- a/spark/common/src/test/java/org/apache/sedona/core/spatialPartitioning/quadtree/QuadTreePartitioningTest.java +++ b/spark/common/src/test/java/org/apache/sedona/core/spatialPartitioning/quadtree/QuadTreePartitioningTest.java @@ -18,15 +18,24 @@ */ package org.apache.sedona.core.spatialPartitioning.quadtree; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + import java.util.ArrayList; +import java.util.HashSet; +import java.util.Iterator; import java.util.List; +import java.util.Set; import org.apache.sedona.core.spatialPartitioning.QuadtreePartitioning; import org.junit.Assert; import org.junit.Test; import org.locationtech.jts.geom.Coordinate; import org.locationtech.jts.geom.Envelope; +import org.locationtech.jts.geom.Geometry; import org.locationtech.jts.geom.GeometryFactory; +import org.locationtech.jts.geom.LineString; import org.locationtech.jts.geom.Point; +import scala.Tuple2; public class QuadTreePartitioningTest { @@ -53,4 +62,44 @@ public void testDataSkew() throws Exception { QuadtreePartitioning partitioning = new QuadtreePartitioning(samples, extent, 10); Assert.assertNotNull(partitioning.getPartitionTree()); } + + @Test + public void testFullGeometryKeysMatchPlacedPartitions() throws Exception { + ExtendedQuadTree source = createExtendedTree(true); + + ExtendedQuadTree queryTree = new ExtendedQuadTree<>(source, true, true); + LineString line = + factory.createLineString(new Coordinate[] {new Coordinate(1, 1), new Coordinate(9, 9)}); + Set placedPartitions = new HashSet<>(); + Iterator> placements = queryTree.placeObject(line); + while (placements.hasNext()) { + placedPartitions.add(placements.next()._1()); + } + + assertTrue(placedPartitions.size() > 1); + assertEquals(placedPartitions, queryTree.getKeys(line)); + } + + @Test + public void testLegacyKeysPreserveHistoricalRouting() throws Exception { + ExtendedQuadTree source = createExtendedTree(false); + Point point = factory.createPoint(new Coordinate(5, 5)); + + assertEquals(source.getQuadTree().getKeys(point), source.getKeys(point)); + + ExtendedQuadTree queryTree = new ExtendedQuadTree<>(source, true, false); + Set expectedKeys = + new HashSet<>(source.getSpatialExpandedBoundaryIndex().query(point.getEnvelopeInternal())); + assertEquals(expectedKeys, queryTree.getKeys(point)); + } + + private ExtendedQuadTree createExtendedTree(boolean useKnnSamples) throws Exception { + ExtendedQuadTree source = new ExtendedQuadTree<>(new Envelope(0, 10, 0, 10), 4); + source.insert(new Envelope(new Coordinate(1, 1))); + source.insert(new Envelope(new Coordinate(1, 9))); + source.insert(new Envelope(new Coordinate(9, 1))); + source.insert(new Envelope(new Coordinate(9, 9))); + source.build(1, false, useKnnSamples); + return source; + } } 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 index b2bc35452a9..9558d2ded6f 100644 --- 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,10 +21,16 @@ import static org.junit.Assert.assertEquals; import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.apache.sedona.core.enums.GridType; import org.apache.sedona.core.enums.IndexType; +import org.apache.sedona.core.spatialPartitioning.QuadTreeRTPartitioner; import org.junit.AfterClass; import org.junit.BeforeClass; import org.junit.Test; +import org.locationtech.jts.geom.Coordinate; +import org.locationtech.jts.geom.GeometryFactory; import org.locationtech.jts.geom.Point; import org.locationtech.jts.index.strtree.STRtree; @@ -138,4 +144,75 @@ public void testBuildQuadtreeIndex() throws Exception { List result = spatialRDD.indexedRDD.take(1).get(0).query(spatialRDD.boundaryEnvelope); } } + + @Test + public void testQuadTreeRTreePartitioningUsesObjectSamples() throws Exception { + GeometryFactory geometryFactory = new GeometryFactory(); + List points = + IntStream.range(0, 128) + .mapToObj(x -> geometryFactory.createPoint(new Coordinate((double) x, 0.0))) + .collect(Collectors.toList()); + PointRDD spatialRDD = new PointRDD(sc.parallelize(points, 8)); + spatialRDD.analyze(); + spatialRDD.setNeighborSampleNumber(points.size()); + spatialRDD.setUseKnnSamples(true); + spatialRDD.setDeduplicateKnnSamples(true); + spatialRDD.spatialPartitioning(GridType.QUADTREE_RTREE, 16); + + QuadTreeRTPartitioner partitioner = (QuadTreeRTPartitioner) spatialRDD.getPartitioner(); + int gridCount = partitioner.getGrids().size(); + for (List expandedGrids : partitioner.getOverlappedGrids().values()) { + assertEquals(gridCount, expandedGrids.size()); + } + } + + @Test + public void testQuadTreeRTreePartitioningFallsBackForInsufficientUniqueSamples() + throws Exception { + GeometryFactory geometryFactory = new GeometryFactory(); + List points = + IntStream.range(0, 128) + .mapToObj(x -> geometryFactory.createPoint(new Coordinate((double) (x / 2), 0.0))) + .collect(Collectors.toList()); + PointRDD spatialRDD = new PointRDD(sc.parallelize(points, 8)); + spatialRDD.analyze(); + spatialRDD.setNeighborSampleNumber(65); + spatialRDD.setUseKnnSamples(true); + spatialRDD.setDeduplicateKnnSamples(true); + spatialRDD.spatialPartitioning(GridType.QUADTREE_RTREE, 16); + + QuadTreeRTPartitioner partitioner = (QuadTreeRTPartitioner) spatialRDD.getPartitioner(); + int gridCount = partitioner.getGrids().size(); + Point object = geometryFactory.createPoint(new Coordinate(63.0, 0.0)); + assertEquals( + gridCount, + partitioner.getSTRForOverlappedGrids().query(object.getEnvelopeInternal()).size()); + for (List expandedGrids : partitioner.getOverlappedGrids().values()) { + assertEquals(1, expandedGrids.size()); + } + } + + @Test + public void testQuadTreeRTreePartitioningFallsBackForInsufficientSamples() throws Exception { + GeometryFactory geometryFactory = new GeometryFactory(); + List points = + IntStream.range(0, 128) + .mapToObj(x -> geometryFactory.createPoint(new Coordinate((double) x, 0.0))) + .collect(Collectors.toList()); + PointRDD spatialRDD = new PointRDD(sc.parallelize(points, 8)); + spatialRDD.analyze(); + spatialRDD.setNeighborSampleNumber(points.size() + 1); + spatialRDD.setUseKnnSamples(true); + spatialRDD.spatialPartitioning(GridType.QUADTREE_RTREE, 16); + + QuadTreeRTPartitioner partitioner = (QuadTreeRTPartitioner) spatialRDD.getPartitioner(); + int gridCount = partitioner.getGrids().size(); + for (List expandedGrids : partitioner.getOverlappedGrids().values()) { + assertEquals(1, expandedGrids.size()); + } + Point object = geometryFactory.createPoint(new Coordinate(63.0, 0.0)); + assertEquals( + gridCount, + partitioner.getSTRForOverlappedGrids().query(object.getEnvelopeInternal()).size()); + } } diff --git a/spark/common/src/test/scala/org/apache/sedona/sql/KnnJoinSuite.scala b/spark/common/src/test/scala/org/apache/sedona/sql/KnnJoinSuite.scala index 4ca17598d30..e6da8714a15 100644 --- a/spark/common/src/test/scala/org/apache/sedona/sql/KnnJoinSuite.scala +++ b/spark/common/src/test/scala/org/apache/sedona/sql/KnnJoinSuite.scala @@ -22,8 +22,10 @@ import org.apache.spark.sql.catalyst.expressions.Literal import org.apache.spark.sql.functions.col import org.apache.spark.sql.sedona_sql.UDT.GeometryUDT import org.apache.spark.sql.sedona_sql.expressions.st_constructors.ST_GeomFromText -import org.apache.spark.sql.sedona_sql.strategy.join.KNNJoinExec -import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType} +import org.apache.spark.sql.sedona_sql.strategy.join.{BroadcastObjectSideKNNJoinExec, BroadcastQuerySideKNNJoinExec, KNNJoinExec, TraitKNNJoinQueryExec} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.types.{IntegerType, LongType, StringType, StructField, StructType} import org.apache.spark.sql.{Column, DataFrame, Row, SparkSession} import org.apache.spark.sql.functions.expr import org.scalatest.matchers.must.Matchers.{be, include} @@ -209,6 +211,34 @@ class KnnJoinSuite extends TestBaseScala with TableDrivenPropertyChecks { "[1,3][1,6][1,13][1,16][2,1][2,5][2,11][2,15][3,3][3,9][3,13][3,19]") } + for (broadcastSide <- Seq("queries", "objects")) { + it(s"KNN Join should apply extra conditions when broadcasting $broadcastSide") { + sparkSession + .range(0, 2) + .selectExpr("id", "ST_Point(id, 0.0) AS geom") + .createOrReplaceTempView("BROADCAST_EXTRA_QUERIES") + sparkSession + .range(0, 2) + .selectExpr("id", "ST_Point(id + 10.0, 0.0) AS geom") + .createOrReplaceTempView("BROADCAST_EXTRA_OBJECTS") + + val hint = if (broadcastSide == "queries") "BROADCAST(q)" else "BROADCAST(o)" + val df = sparkSession.sql(s"""SELECT /*+ $hint */ q.id, o.id + |FROM BROADCAST_EXTRA_QUERIES q + |JOIN BROADCAST_EXTRA_OBJECTS o + |ON ST_KNN(q.geom, o.geom, 1, false) + |AND ST_Distance(q.geom, o.geom) <= 1.0""".stripMargin) + + df.count() should be(0) + val plan = findKNNJoinExec(df) + if (broadcastSide == "queries") { + plan.isInstanceOf[BroadcastQuerySideKNNJoinExec] should be(true) + } else { + plan.isInstanceOf[BroadcastObjectSideKNNJoinExec] should be(true) + } + } + } + it("KNN Join should verify the correct parameter k is passed to the join function") { val df = sparkSession .range(0, 1) @@ -442,6 +472,276 @@ class KnnJoinSuite extends TestBaseScala with TableDrivenPropertyChecks { } } + it("KNN Join should support query-local tie controls") { + val points1 = Seq((0, "POINT(0 0)", "query")) + val points2 = + Seq((0, "POINT(-1 0)", "west"), (1, "POINT(1 0)", "east"), (2, "POINT(10 0)", "far")) + val (df1, df2) = createPointsDataFrames(sparkSession, points1, points2) + df1.createOrReplaceTempView("QUERIES") + df2.createOrReplaceTempView("OBJECTS") + + withExportTies(export = false) { + val df = sparkSession.sql("""SELECT QUERIES.ID AS qid, OBJECTS.ID AS oid + |FROM QUERIES JOIN OBJECTS + |ON ST_KNN(QUERIES.GEOM, OBJECTS.GEOM, 1, false, true)""".stripMargin) + + val result = df.collect().sortBy(_.getInt(1)) + result.mkString should be("[0,0][0,1]") + val plan = findKNNJoinExec(df) + plan.includeTies should be(Some(true)) + plan.exclusive should be(false) + } + + withExportTies(export = true) { + val df = sparkSession.sql("""SELECT QUERIES.ID AS qid, OBJECTS.ID AS oid + |FROM QUERIES JOIN OBJECTS + |ON ST_KNN(QUERIES.GEOM, OBJECTS.GEOM, 1, false, false)""".stripMargin) + + val result = df.collect() + result.length should be(1) + Set(0, 1).contains(result.head.getInt(1)) should be(true) + findKNNJoinExec(df).includeTies should be(Some(false)) + } + } + + it("query-local planar ties should preserve duplicate rows in both broadcast plans") { + sparkSession + .range(0, 1, 1, 1) + .selectExpr("cast(id as int) AS id", "ST_Point(cast(id as double), 0.0) AS geom") + .createOrReplaceTempView("BROADCAST_TIE_QUERIES") + sparkSession + .range(0, 2, 1, 1) + .selectExpr("cast(id as int) AS id", "ST_Point(1.0, cast(id - id as double)) AS geom") + .createOrReplaceTempView("BROADCAST_TIE_OBJECTS") + + for (broadcastSide <- Seq("q", "o")) { + val queryLocal = sparkSession.sql(s"""SELECT /*+ BROADCAST($broadcastSide) */ + |q.ID AS qid, o.ID AS oid + |FROM BROADCAST_TIE_QUERIES q JOIN BROADCAST_TIE_OBJECTS o + |ON ST_KNN(q.GEOM, o.GEOM, 1, false, true)""".stripMargin) + + queryLocal.collect().sortBy(_.getInt(1)).mkString should be("[0,0][0,1]") + val queryLocalPlan = findKNNJoinExec(queryLocal) + if (broadcastSide == "q") { + queryLocalPlan.isInstanceOf[BroadcastQuerySideKNNJoinExec] should be(true) + } else { + queryLocalPlan.isInstanceOf[BroadcastObjectSideKNNJoinExec] should be(true) + } + + withExportTies(export = true) { + val legacy = sparkSession.sql(s"""SELECT /*+ BROADCAST($broadcastSide) */ + |q.ID AS qid, o.ID AS oid + |FROM BROADCAST_TIE_QUERIES q JOIN BROADCAST_TIE_OBJECTS o + |ON ST_KNN(q.GEOM, o.GEOM, 1, false)""".stripMargin) + + val legacyRows = legacy.collect() + legacyRows.length should be(1) + Set(0, 1).contains(legacyRows.head.getInt(1)) should be(true) + } + } + } + + it("query-local no-tie selection should be stable across regular and broadcast plans") { + sparkSession + .range(0, 1, 1, 1) + .selectExpr("cast(id as int) AS id", "ST_Point(cast(id - id as double), 0.0) AS geom") + .createOrReplaceTempView("STABLE_TIE_QUERIES") + sparkSession + .range(0, 2, 1, 1) + .selectExpr("cast(id as int) AS id", "ST_Point(cast(id - id + 1 as double), 0.0) AS geom") + .createOrReplaceTempView("STABLE_TIE_OBJECTS") + + withConf( + Map( + "sedona.join.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.autoBroadcastJoinThreshold" -> "-1")) { + def nearestWithHint(hint: String) = + sparkSession.sql(s"""SELECT $hint q.ID AS qid, o.ID AS oid + |FROM STABLE_TIE_QUERIES q JOIN STABLE_TIE_OBJECTS o + |ON ST_KNN(q.GEOM, o.GEOM, 1, false, false)""".stripMargin) + + val regular = nearestWithHint("") + val broadcastQuery = nearestWithHint("/*+ BROADCAST(q) */") + val broadcastObject = nearestWithHint("/*+ BROADCAST(o) */") + + findKNNJoinExec(regular).isInstanceOf[KNNJoinExec] should be(true) + findKNNJoinExec(broadcastQuery).isInstanceOf[BroadcastQuerySideKNNJoinExec] should be( + true) + findKNNJoinExec(broadcastObject).isInstanceOf[BroadcastObjectSideKNNJoinExec] should be( + true) + + Seq(regular, broadcastQuery, broadcastObject).map( + _.collect().map(_.getInt(1)).toSeq) should be(Seq(Seq(0), Seq(0), Seq(0))) + } + } + + it("query-local geography should retain historical planar tie expansion") { + val points1 = Seq((0, "POINT(0 90)", "query")) + val points2 = Seq( + (0, "POINT(0 -90)", "antipode-a"), + (1, "POINT(100 -90)", "antipode-b"), + (2, "POINT(100 -90)", "antipode-b-duplicate")) + val (df1, df2) = createPointsDataFrames(sparkSession, points1, points2) + df1.createOrReplaceTempView("SPHEROID_TIE_QUERIES") + df2.createOrReplaceTempView("SPHEROID_TIE_OBJECTS") + + val df = sparkSession.sql("""SELECT /*+ BROADCAST(o) */ q.ID AS qid, o.ID AS oid + |FROM SPHEROID_TIE_QUERIES q JOIN SPHEROID_TIE_OBJECTS o + |ON ST_KNN(q.GEOM, o.GEOM, 1, true, true)""".stripMargin) + + val rows = df.collect() + rows.length should be(1) + Set(0, 1, 2).contains(rows.head.getInt(1)) should be(true) + findKNNJoinExec(df).isInstanceOf[BroadcastObjectSideKNNJoinExec] should be(true) + } + + it("KNN Join exclusive mode should promote the nearest non-equal geometry") { + val queries = sparkSession + .range(0, 2) + .selectExpr( + "cast(id as int) AS id", + """CASE WHEN id = 0 + |THEN ST_MakeLine(ST_Point(0.0, 0.0), ST_Point(1.0, 0.0)) + |ELSE ST_MakeLine(ST_Point(10.0, 0.0), ST_Point(11.0, 0.0)) END AS geom""".stripMargin) + .repartition(2) + val objects = sparkSession + .range(0, 2) + .selectExpr( + "cast(id as int) AS id", + """CASE WHEN id = 0 + |THEN ST_GeomFromText('LINESTRING(1 0, 0 0)') + |ELSE ST_GeomFromText('LINESTRING(0 2, 1 2)') END AS geom""".stripMargin) + .repartition(2) + queries.createOrReplaceTempView("QUERIES") + objects.createOrReplaceTempView("OBJECTS") + + val df = sparkSession.sql("""SELECT QUERIES.ID AS qid, OBJECTS.ID AS oid + |FROM QUERIES JOIN OBJECTS + |ON ST_KNN(QUERIES.GEOM, OBJECTS.GEOM, 1, false, true, true)""".stripMargin) + + df.where("qid = 0").collect().mkString should be("[0,1]") + val plan = findKNNJoinExec(df) + plan.includeTies should be(Some(true)) + plan.exclusive should be(true) + } + + it("KNN Join exclusive mode should find a distant candidate across partitions") { + withConf( + Map( + "sedona.join.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.autoBroadcastJoinThreshold" -> "-1")) { + val queries = sparkSession + .range(0, 2) + .selectExpr( + "cast(id as int) AS id", + "CASE WHEN id = 0 THEN ST_Point(0.0, 0.0) ELSE ST_Point(2000.0, 0.0) END AS geom") + .repartition(4) + val objects = sparkSession + .range(0, 101) + .selectExpr( + "cast(id as int) AS id", + "CASE WHEN id < 100 THEN ST_Point(0.0, 0.0) ELSE ST_Point(1000.0, 0.0) END AS geom") + .repartition(8) + queries.createOrReplaceTempView("QUERIES") + objects.createOrReplaceTempView("OBJECTS") + + val df = sparkSession.sql("""SELECT QUERIES.ID AS qid, OBJECTS.ID AS oid + |FROM QUERIES JOIN OBJECTS + |ON ST_KNN(QUERIES.GEOM, OBJECTS.GEOM, 1, false, true, true)""".stripMargin) + + df.where("qid = 0").collect().mkString should be("[0,100]") + findKNNJoinExec(df).exclusive should be(true) + } + } + + it("query-local KNN should use a line's full geometry in the regular plan") { + withConf( + Map( + "sedona.join.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.autoBroadcastJoinThreshold" -> "-1")) { + val queries = + createGeometryDataFrame(Seq(0 -> "LINESTRING (0 0, 1000000000 0)"), 64) + val objects = sparkSession + .range(0, 10001, 1, 64) + .selectExpr( + "cast(id as int) AS id", + """CASE WHEN id = 0 THEN ST_Point(0.0, 0.0) + |ELSE ST_Point(500000000.0 + cast(id as double), 10.0) END AS geom""".stripMargin) + + val df = queries + .alias("q") + .join(objects.alias("o"), expr("ST_KNN(q.geom, o.geom, 1, false, false)")) + .selectExpr("q.id AS qid", "o.id AS oid") + + findKNNJoinExec(df).isInstanceOf[KNNJoinExec] should be(true) + df.collect().mkString should be("[0,0]") + } + } + + it("query-local KNN should use a polygon's full geometry in the regular plan") { + withConf( + Map( + "sedona.join.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.autoBroadcastJoinThreshold" -> "-1")) { + val queries = createGeometryDataFrame( + Seq(0 -> "POLYGON ((0 0, 1000000000 0, 1000000000 100, 0 100, 0 0))"), + 32) + val objects = sparkSession + .range(0, 4097, 1, 32) + .selectExpr( + "cast(id as int) AS id", + """CASE WHEN id = 0 THEN ST_Point(0.0, 50.0) + |ELSE ST_Point(500000000.0 + cast(id as double), 110.0) END AS geom""".stripMargin) + + val df = queries + .alias("q") + .join(objects.alias("o"), expr("ST_KNN(q.geom, o.geom, 1, false, false)")) + .selectExpr("q.id AS qid", "o.id AS oid") + + findKNNJoinExec(df).isInstanceOf[KNNJoinExec] should be(true) + df.collect().mkString should be("[0,0]") + } + } + + it("query-local regular KNN should preserve duplicate spanning rows without replicas") { + withConf( + Map( + "sedona.join.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.autoBroadcastJoinThreshold" -> "-1")) { + val queries = createGeometryDataFrame( + Seq(0 -> "LINESTRING (0 0, 1000000000 0)", 1 -> "LINESTRING (0 0, 1000000000 0)"), + 32) + val objects = sparkSession + .range(0, 4098, 1, 32) + .selectExpr( + "cast(id as int) AS id", + """CASE WHEN id < 2 + |THEN ST_GeomFromText('MULTIPOINT ((0 0), (1000000000 0))') + |ELSE ST_Point(500000000.0 + cast(id as double), 10.0) END AS geom""".stripMargin) + + val includeTiesResult = queries + .alias("q") + .join(objects.alias("o"), expr("ST_KNN(q.geom, o.geom, 1, false, true)")) + .selectExpr("q.id AS qid", "o.id AS oid") + findKNNJoinExec(includeTiesResult).isInstanceOf[KNNJoinExec] should be(true) + includeTiesResult + .collect() + .sortBy(row => (row.getInt(0), row.getInt(1))) + .mkString should be("[0,0][0,1][1,0][1,1]") + + val excludeTiesResult = queries + .alias("q") + .join(objects.alias("o"), expr("ST_KNN(q.geom, o.geom, 1, false, false)")) + .selectExpr("q.id AS qid", "o.id AS oid") + val untiedRows = excludeTiesResult.collect() + untiedRows.length should be(2) + untiedRows + .groupBy(_.getInt(0)) + .map { case (queryId, rows) => queryId -> rows.length } should be(Map(0 -> 1, 1 -> 1)) + untiedRows.foreach(row => Set(0, 1).contains(row.getInt(1)) should be(true)) + } + } + it("KNN Join with exact algorithms should not fail with null geometries") { val df1 = sparkSession.sql( "SELECT ST_GeomFromText(col1) as geom1 from values ('POINT (0.0 0.0)'), (null)") @@ -459,6 +759,42 @@ class KnnJoinSuite extends TestBaseScala with TableDrivenPropertyChecks { df1.join(df2, expr("ST_KNN(geom1, geom2, 1)")).count() should be(0) } + it("KNN Join should return no matches for empty or all-null candidate sides") { + val queries = sparkSession + .range(0, 2) + .selectExpr("id", "ST_Point(id, 0.0) AS geom") + val emptyObjects = sparkSession + .createDataFrame( + sparkSession.sparkContext.emptyRDD[Row], + StructType( + Seq( + StructField("id", LongType, nullable = false), + StructField("wkt", StringType, nullable = true)))) + .selectExpr("id", "ST_GeomFromText(wkt) AS geom") + val nullObjects = sparkSession + .createDataFrame( + sparkSession.sparkContext.parallelize(Seq(Row(0L, null), Row(1L, null)), 2), + StructType( + Seq( + StructField("id", LongType, nullable = false), + StructField("wkt", StringType, nullable = true)))) + .selectExpr("id", "ST_GeomFromText(wkt) AS geom") + + val emptyResult = queries + .alias("queries") + .join( + emptyObjects.alias("emptyObjects"), + expr("ST_KNN(queries.geom, emptyObjects.geom, 1)")) + val nullResult = queries + .alias("queries") + .join(nullObjects.alias("nullObjects"), expr("ST_KNN(queries.geom, nullObjects.geom, 1)")) + + emptyResult.queryExecution.executedPlan.toString should include("KNNJoin") + nullResult.queryExecution.executedPlan.toString should include("KNNJoin") + emptyResult.count() should be(0) + nullResult.count() should be(0) + } + it("KNN Join using spider data source") { val dfRandomSquares = sparkSession.read .format("spider") @@ -504,6 +840,18 @@ class KnnJoinSuite extends TestBaseScala with TableDrivenPropertyChecks { withConf(Map("spark.sedona.join.debug.spatialPartitionerSavePath" -> path))(body) } + private def findKNNJoinExec(df: DataFrame): TraitKNNJoinQueryExec = { + def collect(plan: SparkPlan): Option[TraitKNNJoinQueryExec] = + plan.collectFirst { case knn: TraitKNNJoinQueryExec => knn } + + val executedPlan = df.queryExecution.executedPlan + collect(executedPlan) + .orElse(executedPlan.collectFirst { case adaptive: AdaptiveSparkPlanExec => + collect(adaptive.executedPlan) + }.flatten) + .get + } + def validateQueryPlan( df: DataFrame, numNeighbors: Int, @@ -681,4 +1029,17 @@ class KnnJoinSuite extends TestBaseScala with TableDrivenPropertyChecks { (df1, df2) } + + private def createGeometryDataFrame( + geometries: Seq[(Int, String)], + partitions: Int): DataFrame = { + val schema = StructType( + Seq( + StructField("id", IntegerType, nullable = false), + StructField("wkt", StringType, nullable = false))) + val rows = geometries.map { case (id, wkt) => Row(id, wkt) } + sparkSession + .createDataFrame(sparkSession.sparkContext.parallelize(rows, partitions), schema) + .selectExpr("id", "ST_GeomFromText(wkt) AS geom") + } }