Add support for ST_DWithin in SPATIAL_JOIN

src/spatial/modules/main/spatial_functions_scalar.cpp

@@ -16,6 +16,8 @@
 #include "duckdb/planner/expression/bound_function_expression.hpp"
 #include "duckdb/common/vector_operations/septenary_executor.hpp"
 
+#include "spatial/util/distance_extract.hpp"
+
 // Extra
 #include "yyjson.h"
 
@@ -8562,6 +8564,18 @@ struct ST_MMin : VertexAggFunctionBase<ST_MMin, VertexMinAggOp> {
 
 } // namespace
 
+// Helper to access the constant distance from the bind data
+bool ST_DWithinHelper::TryGetConstDistance(const unique_ptr<FunctionData> &bind_data, double &result) {
+	if (bind_data) {
+		const auto &data = bind_data->Cast<ST_DistanceWithin::BindData>();
+		if (data.is_constant) {
+			result = data.distance;
+			return true;
+		}
+	}
+	return false;
+}
+
 //######################################################################################################################
 // Register
 //######################################################################################################################

src/spatial/operators/spatial_join_logical.cpp

@@ -76,7 +76,7 @@ PhysicalOperator &LogicalSpatialJoin::CreatePlan(ClientContext &context, Physica
 	auto &right = generator.CreatePlan(*children[1]);
 
 	return generator.Make<PhysicalSpatialJoin>(*this, left, right, std::move(spatial_predicate), join_type,
-	                                           estimated_cardinality);
+	                                           estimated_cardinality, has_const_distance, const_distance);
 }
 
 void LogicalSpatialJoin::Serialize(Serializer &writer) const {
@@ -89,6 +89,10 @@ void LogicalSpatialJoin::Serialize(Serializer &writer) const {
 	writer.WritePropertyWithDefault<vector<idx_t>>(403, "right_projection_map", right_projection_map);
 	writer.WritePropertyWithDefault<unique_ptr<Expression>>(404, "spatial_predicate", spatial_predicate);
 	writer.WritePropertyWithDefault<vector<unique_ptr<Expression>>>(405, "extra_conditions", extra_conditions);
+	writer.WritePropertyWithDefault<bool>(406, "has_const_distance", has_const_distance);
+	if (has_const_distance) {
+		writer.WritePropertyWithDefault<double>(407, "const_distance", const_distance);
+	}
 }
 
 unique_ptr<LogicalExtensionOperator> LogicalSpatialJoin::Deserialize(Deserializer &reader) {
@@ -98,13 +102,22 @@ unique_ptr<LogicalExtensionOperator> LogicalSpatialJoin::Deserialize(Deserialize
 	auto right_projection_map = reader.ReadPropertyWithDefault<vector<idx_t>>(403, "right_projection_map");
 	auto spatial_predicate = reader.ReadPropertyWithDefault<unique_ptr<Expression>>(404, "spatial_predicate");
 	auto extra_conditions = reader.ReadPropertyWithDefault<vector<unique_ptr<Expression>>>(405, "extra_conditions");
+	auto has_const_distance = reader.ReadPropertyWithExplicitDefault<bool>(406, "has_const_distance", false);
+	double const_distance = 0.0;
+	if (has_const_distance) {
+		const_distance = reader.ReadPropertyWithDefault<double>(407, "const_distance");
+	}
 
 	auto result = make_uniq<LogicalSpatialJoin>(join_type);
 	result->mark_index = mark_index;
 	result->left_projection_map = std::move(left_projection_map);
 	result->right_projection_map = std::move(right_projection_map);
 	result->spatial_predicate = std::move(spatial_predicate);
 	result->extra_conditions = std::move(extra_conditions);
+	result->has_const_distance = has_const_distance;
+	if (has_const_distance) {
+		result->const_distance = const_distance;
+	}
 
 	return std::move(result);
 }

src/spatial/operators/spatial_join_logical.hpp

@@ -28,6 +28,9 @@ class LogicalSpatialJoin final : public LogicalExtensionOperator {
 	//! Join Keys statistics (optional)
 	vector<unique_ptr<BaseStatistics>> join_stats;
 
+	bool has_const_distance = false;
+	double const_distance = 0.0;
+
 public:
 	explicit LogicalSpatialJoin(JoinType join_type_p);
 

src/spatial/operators/spatial_join_optimizer.cpp

@@ -1,19 +1,20 @@
 #include "spatial_join_optimizer.hpp"
+#include "spatial_join_logical.hpp"
+#include "spatial/util/distance_extract.hpp"
 
 #include "duckdb/main/database.hpp"
 #include "duckdb/planner/expression/bound_function_expression.hpp"
 #include "duckdb/planner/operator/logical_any_join.hpp"
-#include "spatial_join_logical.hpp"
-
 #include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp"
 #include "duckdb/planner/operator/logical_filter.hpp"
 
 namespace duckdb {
 
 // All of these imply bounding box intersection
 static case_insensitive_set_t spatial_predicate_map = {
-    "ST_Equals",   "ST_Intersects", "ST_Touches",   "ST_Crosses",          "ST_Within",        "ST_Contains",
-    "ST_Overlaps", "ST_Covers",     "ST_CoveredBy", "ST_ContainsProperly", "ST_WithinProperly"};
+    "ST_Equals",   "ST_Intersects", "ST_Touches",   "ST_Crosses",          "ST_Within",         "ST_Contains",
+    "ST_Overlaps", "ST_Covers",     "ST_CoveredBy", "ST_ContainsProperly", "ST_WithinProperly", "ST_DWithin",
+};
 
 static case_insensitive_map_t<string> spatial_predicate_inverse_map = {
     {"ST_Equals", "ST_Equals"},
@@ -26,7 +27,9 @@ static case_insensitive_map_t<string> spatial_predicate_inverse_map = {
     {"ST_Covers", "ST_CoveredBy"},                // Inverse
     {"ST_CoveredBy", "ST_Covers"},                // Inverse
     {"ST_WithinProperly", "ST_ContainsProperly"}, // Inverse
-    {"ST_ContainsProperly", "ST_WithinProperly"}  // Inverse
+    {"ST_ContainsProperly", "ST_WithinProperly"}, // Inverse
+    {"ST_DWithin", "ST_DWithin"},                 // Symmetric (when distance is constant)
+
 };
 
 unique_ptr<Expression> TryGetInversePredicate(ClientContext &context, unique_ptr<Expression> expr) {
@@ -130,6 +133,12 @@ static void InsertSpatialJoin(OptimizerExtensionInput &input, unique_ptr<Logical
 
 		auto &func = expr->Cast<BoundFunctionExpression>();
 
+		// The function must be a binary predicate
+		if (func.children.size() != 2) {
+			extra_predicates.push_back(std::move(expr));
+			continue;
+		}
+
 		// The function must be a recognized spatial predicate
 		if (spatial_predicate_map.count(func.function.name) == 0) {
 			extra_predicates.push_back(std::move(expr));
@@ -179,6 +188,14 @@ static void InsertSpatialJoin(OptimizerExtensionInput &input, unique_ptr<Logical
 	spatial_join->has_estimated_cardinality = any_join.has_estimated_cardinality;
 	spatial_join->estimated_cardinality = any_join.estimated_cardinality;
 
+	// If this is ST_DWithin, try to extract the constant distance value
+	const auto &pred_func = spatial_join->spatial_predicate->Cast<BoundFunctionExpression>();
+	if (pred_func.function.name == "ST_DWithin") {
+		// Try to get the constant distance value from the bind data;
+		spatial_join->has_const_distance =
+		    ST_DWithinHelper::TryGetConstDistance(pred_func.bind_info, spatial_join->const_distance);
+	}
+
 	// Replace the operator
 	plan = std::move(spatial_join);
 }

src/spatial/operators/spatial_join_physical.cpp

@@ -330,9 +330,9 @@ class FlatRTree {
 
 PhysicalSpatialJoin::PhysicalSpatialJoin(LogicalOperator &op, PhysicalOperator &left, PhysicalOperator &right,
                                          unique_ptr<Expression> condition_p, JoinType join_type,
-                                         idx_t estimated_cardinality)
+                                         idx_t estimated_cardinality, bool has_const_distance, double const_distance)
     : PhysicalJoin(op, PhysicalOperatorType::EXTENSION, join_type, estimated_cardinality),
-      condition(std::move(condition_p)) {
+      condition(std::move(condition_p)), has_const_distance(has_const_distance), const_distance(const_distance) {
 
 	children.emplace_back(left);
 	children.emplace_back(right);
@@ -636,6 +636,15 @@ SinkFinalizeType PhysicalSpatialJoin::Finalize(Pipeline &pipeline, Event &event,
 				continue;
 			}
 
+			if (has_const_distance) {
+				// If this is a ST_DWithin join, we need to expand the bounding box by the constant distance
+				const auto f_dist = MathUtil::DoubleToFloatUp(const_distance);
+				bbox.min.x -= f_dist;
+				bbox.min.y -= f_dist;
+				bbox.max.x += f_dist;
+				bbox.max.y += f_dist;
+			}
+
 			// Push the bounding box into the R-Tree
 			gstate.rtree->Push(bbox, rows_ptr[row_idx]);
 		}

src/spatial/operators/spatial_join_physical.hpp

@@ -11,7 +11,8 @@ class PhysicalSpatialJoin final : public PhysicalJoin {
 
 public:
 	PhysicalSpatialJoin(LogicalOperator &op, PhysicalOperator &left, PhysicalOperator &right,
-	                    unique_ptr<Expression> spatial_predicate, JoinType join_type, idx_t estimated_cardinality);
+	                    unique_ptr<Expression> spatial_predicate, JoinType join_type, idx_t estimated_cardinality,
+	                    bool has_const_distance, double const_distance);
 
 	//! The condition of the join
 	unique_ptr<Expression> condition;
@@ -31,6 +32,10 @@ class PhysicalSpatialJoin final : public PhysicalJoin {
 	shared_ptr<TupleDataLayout> layout;
 	idx_t build_side_match_offset = 0; // This is the byte offset to the match column for right/outer joins
 
+	// In case this is a ST_DWithin join, we store the constant distance here
+	bool has_const_distance = false;
+	double const_distance = 0.0;
+
 public:
 	// Operator Interface
 	unique_ptr<OperatorState> GetOperatorState(ExecutionContext &context) const override;

src/spatial/util/distance_extract.hpp

@@ -0,0 +1,12 @@
+#pragma once
+
+namespace duckdb {
+
+struct FunctionData;
+
+// Helper class to extract the constant distance from the ST_DWithin function.
+struct ST_DWithinHelper {
+	static bool TryGetConstDistance(const unique_ptr<FunctionData> &bind_data, double &result);
+};
+
+} // namespace duckdb