feat: add ST_KNN spatial join operator (in-memory, planar)

Adds a K-nearest neighbor spatial join to duckdb-spatial:

- ST_KNN(geom1, geom2, k) — scalar stub recognized by the optimizer
- Optimizer rewrites JOIN ON ST_KNN(...) into SPATIAL_KNN_JOIN
- Hjaltason-Samet priority queue KNN on FlatRTree
- Exact distance refinement with adaptive overfetch (2x, retry 8x)
- INNER and LEFT JOIN support
- JOO child swap handling via build_child_idx

Stripped for review (follow-up PRs):
- No out-of-core / spill-to-disk (tree always in memory)
- No spheroid support (planar distance only)
- No partitioning (single R-tree)

The FlatRTree is compact (~24 bytes/entry), so in-memory handles
up to ~100M entries (2.4GB tree) on commodity hardware.

Author: pierre-warnier

src/spatial/modules/main/spatial_functions_scalar.cpp

@@ -18,6 +18,7 @@
 #include "duckdb/planner/expression/bound_constant_expression.hpp"
 
 #include "spatial/util/distance_extract.hpp"
+#include "spatial/util/knn_extract.hpp"
 #include "spatial/spatial_settings.hpp"
 
 // Extra
@@ -9462,6 +9463,76 @@ constexpr const char *ST_X::NAME;
 constexpr const char *ST_Y::NAME;
 constexpr const char *ST_Z::NAME;
 
+//======================================================================================================================
+// ST_KNN
+//======================================================================================================================
+
+struct ST_KNN {
+
+	class BindData final : public FunctionData {
+	public:
+		int32_t k;
+		bool is_constant = false;
+
+		BindData(int32_t k_p) : k(k_p), is_constant(true) {
+		}
+
+		unique_ptr<FunctionData> Copy() const override {
+			return make_uniq<BindData>(k);
+		}
+
+		bool Equals(const FunctionData &other) const override {
+			auto &other_data = other.Cast<BindData>();
+			return is_constant == other_data.is_constant && k == other_data.k;
+		}
+	};
+
+	static unique_ptr<FunctionData> Bind3(ClientContext &context, ScalarFunction &bound_function,
+	                                      vector<unique_ptr<Expression>> &arguments) {
+		// ST_KNN(geom1, geom2, k)
+		if (arguments[2]->IsFoldable()) {
+			const auto k_expr = ExpressionExecutor::EvaluateScalar(context, *arguments[2]);
+			const auto k_value = k_expr.GetValue<int32_t>();
+			if (k_value < 1) {
+				throw InvalidInputException("ST_KNN: k must be >= 1, got %d", k_value);
+			}
+			Function::EraseArgument(bound_function, arguments, 2);
+			return make_uniq<BindData>(k_value);
+		}
+		throw InvalidInputException("ST_KNN: k must be a constant expression");
+	}
+
+	static void Execute(DataChunk &args, ExpressionState &state, Vector &result) {
+		throw InvalidInputException("ST_KNN cannot be used outside of a JOIN ON clause");
+	}
+
+	static void Register(ExtensionLoader &loader) {
+		FunctionBuilder::RegisterScalar(loader, "ST_KNN", [](ScalarFunctionBuilder &func) {
+			// ST_KNN(geom1, geom2, k)
+			func.AddVariant([](ScalarFunctionVariantBuilder &variant) {
+				variant.AddParameter("geom1", LogicalType::GEOMETRY());
+				variant.AddParameter("geom2", LogicalType::GEOMETRY());
+				variant.AddParameter("k", LogicalType::INTEGER);
+				variant.SetReturnType(LogicalType::BOOLEAN);
+				variant.SetBind(Bind3);
+				variant.SetFunction(Execute);
+			});
+			func.SetDescription(R"(
+				K-nearest neighbor spatial join predicate.
+				Finds the k nearest geometries from geom2 for each geom1.
+				Must be used in a JOIN ON clause.
+			)");
+			func.SetExample(R"(
+				SELECT a.id, b.id
+				FROM table_a a
+				JOIN table_b b ON ST_KNN(a.geom, b.geom, 5);
+			)");
+			func.SetTag("ext", "spatial");
+			func.SetTag("category", "relation");
+		});
+	}
+};
+
 } // namespace
 
 // Helper to access the constant distance from the bind data
@@ -9476,6 +9547,19 @@ bool ST_DWithinHelper::TryGetConstDistance(const unique_ptr<FunctionData> &bind_
 	return false;
 }
 
+// Helper to access the constant k from the bind data
+bool ST_KNNHelper::TryGetConstK(const unique_ptr<FunctionData> &bind_data, int32_t &result) {
+	if (bind_data) {
+		const auto &data = bind_data->Cast<ST_KNN::BindData>();
+		if (data.is_constant) {
+			result = data.k;
+			return true;
+		}
+	}
+	return false;
+}
+
+
 //######################################################################################################################
 // Register
 //######################################################################################################################
@@ -9528,6 +9612,7 @@ void RegisterSpatialScalarFunctions(ExtensionLoader &loader) {
 	ST_InterpolatePoint::Register(loader);
 	ST_Intersects::Register(loader);
 	ST_Intersects_Extent::Register(loader);
+	ST_KNN::Register(loader);
 	ST_IsClosed::Register(loader);
 	ST_IsEmpty::Register(loader);
 	ST_Length::Register(loader);

src/spatial/operators/CMakeLists.txt

@@ -3,5 +3,7 @@ set(EXTENSION_SOURCES
     ${CMAKE_CURRENT_SOURCE_DIR}/spatial_join_logical.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/spatial_join_physical.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/spatial_join_optimizer.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/spatial_knn_join_logical.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/spatial_knn_join_physical.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/spatial_operator_extension.cpp
     PARENT_SCOPE)

src/spatial/operators/spatial_join_optimizer.cpp

@@ -1,6 +1,8 @@
 #include "spatial_join_optimizer.hpp"
 #include "spatial_join_logical.hpp"
+#include "spatial/operators/spatial_knn_join_logical.hpp"
 #include "spatial/util/distance_extract.hpp"
+#include "spatial/util/knn_extract.hpp"
 #include "spatial/spatial_types.hpp"
 
 #include "duckdb/main/database.hpp"
@@ -319,7 +321,168 @@ static void TrySwapAnyJoin(OptimizerExtensionInput &input, unique_ptr<LogicalOpe
 	plan = std::move(spatial_join);
 }
 
+//======================================================================================================================
+// KNN Join Detection
+//======================================================================================================================
+
+static bool IsKNNJoinPredicate(const unique_ptr<Expression> &expr, const unordered_set<idx_t> &left_bindings,
+                               const unordered_set<idx_t> &right_bindings, bool &needs_flipping, int32_t &k_value) {
+
+	const auto total_side = JoinSide::GetJoinSide(*expr, left_bindings, right_bindings);
+	if (total_side != JoinSide::BOTH) {
+		return false;
+	}
+
+	if (expr->type != ExpressionType::BOUND_FUNCTION) {
+		return false;
+	}
+
+	auto &func = expr->Cast<BoundFunctionExpression>();
+
+	if (!StringUtil::CIEquals(func.function.name, "ST_KNN")) {
+		return false;
+	}
+
+	// After bind, ST_KNN has 2 args (k was folded into bind_data)
+	if (func.children.size() != 2) {
+		return false;
+	}
+
+	if (func.return_type != LogicalType::BOOLEAN) {
+		return false;
+	}
+
+	const auto left_side = JoinSide::GetJoinSide(*func.children[0], left_bindings, right_bindings);
+	const auto right_side = JoinSide::GetJoinSide(*func.children[1], left_bindings, right_bindings);
+
+	if (left_side == JoinSide::BOTH || right_side == JoinSide::BOTH) {
+		return false;
+	}
+
+	needs_flipping = (left_side == JoinSide::RIGHT);
+
+	if (!ST_KNNHelper::TryGetConstK(func.bind_info, k_value)) {
+		return false;
+	}
+
+	return true;
+}
+
+static bool TrySwapKNNAnyJoin(OptimizerExtensionInput &input, unique_ptr<LogicalOperator> &plan) {
+	auto &op = *plan;
+
+	if (op.type != LogicalOperatorType::LOGICAL_ANY_JOIN) {
+		return false;
+	}
+
+	auto &any_join = op.Cast<LogicalAnyJoin>();
+
+	// KNN supports INNER and LEFT joins.
+	// RIGHT may appear when JOO swaps a LEFT join — we convert it back to LEFT after child swap.
+	// Explicit RIGHT joins from the user are not supported (would need build-side outer emit).
+	if (any_join.join_type != JoinType::INNER && any_join.join_type != JoinType::LEFT &&
+	    any_join.join_type != JoinType::RIGHT) {
+		return false;
+	}
+
+	auto &left_child = any_join.children[0];
+	auto &right_child = any_join.children[1];
+	unordered_set<idx_t> left_bindings;
+	unordered_set<idx_t> right_bindings;
+	LogicalJoin::GetTableReferences(*left_child, left_bindings);
+	LogicalJoin::GetTableReferences(*right_child, right_bindings);
+
+	// Split the join condition by AND
+	vector<unique_ptr<Expression>> expressions;
+	expressions.push_back(any_join.condition->Copy());
+	LogicalFilter::SplitPredicates(expressions);
+
+	unique_ptr<Expression> knn_pred_expr = nullptr;
+	vector<unique_ptr<Expression>> extra_predicates;
+	int32_t k_value = 1;
+
+	for (auto &expr : expressions) {
+		bool unused_flip = false;
+		int32_t k_tmp = 0;
+		if (!knn_pred_expr && IsKNNJoinPredicate(expr, left_bindings, right_bindings, unused_flip, k_tmp)) {
+			knn_pred_expr = std::move(expr);
+			k_value = k_tmp;
+		} else if (expr) {
+			extra_predicates.push_back(std::move(expr));
+		}
+	}
+
+	if (!knn_pred_expr) {
+		return false;
+	}
+
+	// For non-INNER joins, extra predicates can't be pushed as filters safely
+	if (!extra_predicates.empty() && any_join.join_type != JoinType::INNER) {
+		return false;
+	}
+
+	// ST_KNN(probe_geom, build_geom, k): arg0 = probe side, arg1 = build side.
+	// DuckDB sinks children[1], so we need the build side at children[1].
+	// Determine which child the build key (second arg) currently references.
+	auto &knn_func = knn_pred_expr->Cast<BoundFunctionExpression>();
+	auto build_key_side = JoinSide::GetJoinSide(*knn_func.children[1], left_bindings, right_bindings);
+	bool needs_child_swap = (build_key_side == JoinSide::LEFT);
+
+
+	// When JOO swaps children, LEFT becomes RIGHT. Convert back to LEFT after our swap.
+	// If RIGHT wasn't produced by JOO swap, reject it — we don't support explicit RIGHT.
+	auto effective_join_type = any_join.join_type;
+	if (effective_join_type == JoinType::RIGHT) {
+		if (needs_child_swap) {
+			effective_join_type = JoinType::LEFT;
+		} else {
+			// Explicit RIGHT JOIN from user — not supported for KNN
+			return false;
+		}
+	}
+
+	auto knn_join = make_uniq<LogicalSpatialKNNJoin>(effective_join_type);
+	knn_join->spatial_predicate = std::move(knn_pred_expr);
+	knn_join->k = k_value;
+	knn_join->children = std::move(any_join.children);
+	knn_join->expressions = std::move(any_join.expressions);
+
+	if (needs_child_swap) {
+		std::swap(knn_join->children[0], knn_join->children[1]);
+		knn_join->left_projection_map = std::move(any_join.right_projection_map);
+		knn_join->right_projection_map = std::move(any_join.left_projection_map);
+	} else {
+		knn_join->left_projection_map = std::move(any_join.left_projection_map);
+		knn_join->right_projection_map = std::move(any_join.right_projection_map);
+	}
+
+	// build_child_idx is always 1 after potential swap
+	knn_join->build_child_idx = 1;
+
+	knn_join->join_stats = std::move(any_join.join_stats);
+	knn_join->has_estimated_cardinality = any_join.has_estimated_cardinality;
+	knn_join->estimated_cardinality = any_join.estimated_cardinality;
+
+	if (!extra_predicates.empty()) {
+		// Wrap KNN join in a filter for extra AND predicates (INNER only)
+		auto filter = make_uniq<LogicalFilter>();
+		filter->expressions = std::move(extra_predicates);
+		filter->children.push_back(std::move(knn_join));
+		plan = std::move(filter);
+	} else {
+		plan = std::move(knn_join);
+	}
+	return true;
+}
+
+//======================================================================================================================
+
 static void InsertSpatialJoin(OptimizerExtensionInput &input, unique_ptr<LogicalOperator> &plan) {
+	// Try KNN first (more specific)
+	if (TrySwapKNNAnyJoin(input, plan)) {
+		return;
+	}
+
 	if (TrySwapComparisonJoin(input, plan)) {
 		return;
 	}

src/spatial/operators/spatial_knn_join_logical.cpp

@@ -0,0 +1,77 @@
+#include "spatial/operators/spatial_knn_join_logical.hpp"
+#include "spatial/operators/spatial_knn_join_physical.hpp"
+
+#include "duckdb/planner/expression/bound_function_expression.hpp"
+#include "duckdb/execution/column_binding_resolver.hpp"
+#include "duckdb/common/serializer/serializer.hpp"
+#include "duckdb/common/serializer/deserializer.hpp"
+
+namespace duckdb {
+
+LogicalSpatialKNNJoin::LogicalSpatialKNNJoin(JoinType join_type_p) : join_type(join_type_p) {
+}
+
+vector<ColumnBinding> LogicalSpatialKNNJoin::GetColumnBindings() {
+	auto left_bindings = MapBindings(children[0]->GetColumnBindings(), left_projection_map);
+	auto right_bindings = MapBindings(children[1]->GetColumnBindings(), right_projection_map);
+	left_bindings.insert(left_bindings.end(), right_bindings.begin(), right_bindings.end());
+	return left_bindings;
+}
+
+void LogicalSpatialKNNJoin::ResolveColumnBindings(ColumnBindingResolver &res, vector<ColumnBinding> &bindings) {
+	auto &cond = spatial_predicate->Cast<BoundFunctionExpression>();
+
+	// After the optimizer, children are always in correct order:
+	// children[0] = probe (arg0), children[1] = build (arg1).
+	res.VisitOperator(*children[0]);
+	res.VisitExpression(&cond.children[0]);
+	res.VisitOperator(*children[1]);
+	res.VisitExpression(&cond.children[1]);
+
+	bindings = GetColumnBindings();
+}
+
+void LogicalSpatialKNNJoin::ResolveTypes() {
+	types = MapTypes(children[0]->types, left_projection_map);
+	auto right_types = MapTypes(children[1]->types, right_projection_map);
+	types.insert(types.end(), right_types.begin(), right_types.end());
+}
+
+PhysicalOperator &LogicalSpatialKNNJoin::CreatePlan(ClientContext &context, PhysicalPlanGenerator &generator) {
+	auto &left = generator.CreatePlan(*children[0]);
+	auto &right = generator.CreatePlan(*children[1]);
+
+	return generator.Make<PhysicalSpatialKNNJoin>(*this, left, right, std::move(spatial_predicate), join_type,
+	                                              estimated_cardinality, k);
+}
+
+void LogicalSpatialKNNJoin::Serialize(Serializer &writer) const {
+	LogicalExtensionOperator::Serialize(writer);
+	writer.WritePropertyWithDefault(300, "operator_type", string(OPERATOR_TYPE_NAME));
+	writer.WritePropertyWithDefault<JoinType>(400, "join_type", join_type, JoinType::INNER);
+	writer.WritePropertyWithDefault<vector<idx_t>>(402, "left_projection_map", left_projection_map);
+	writer.WritePropertyWithDefault<vector<idx_t>>(403, "right_projection_map", right_projection_map);
+	writer.WritePropertyWithDefault<unique_ptr<Expression>>(404, "spatial_predicate", spatial_predicate);
+	writer.WritePropertyWithDefault<int32_t>(405, "k", k, 1);
+	writer.WritePropertyWithDefault<idx_t>(406, "build_child_idx", build_child_idx, static_cast<idx_t>(1));
+}
+
+unique_ptr<LogicalExtensionOperator> LogicalSpatialKNNJoin::Deserialize(Deserializer &reader) {
+	auto join_type = reader.ReadPropertyWithExplicitDefault<JoinType>(400, "join_type", JoinType::INNER);
+	auto left_projection_map = reader.ReadPropertyWithDefault<vector<idx_t>>(402, "left_projection_map");
+	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 k = reader.ReadPropertyWithExplicitDefault<int32_t>(405, "k", 1);
+	auto build_child_idx = reader.ReadPropertyWithExplicitDefault<idx_t>(406, "build_child_idx", static_cast<idx_t>(1));
+
+	auto result = make_uniq<LogicalSpatialKNNJoin>(join_type);
+	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->k = k;
+	result->build_child_idx = build_child_idx;
+
+	return std::move(result);
+}
+
+} // namespace duckdb