feat: support ST_DWithin pushdown in vortex (#8625)

## Summary

<!--
Why are you proposing this change, and what is its impact? 
Is it part of a long term effort, or a bigger change?

If this PR is related to a tracked effort or an open issue, please link
to the relevant issue.
-->

Insert non-throwing geo predicate `vortex_dwithin` in DuckDB, which
later pushdown into vortex, call `distance` scalar function on scanning,
significantly improve Q1/Q3 performance in SpatialBench.

## What changes are included in this PR?
<!--
No need to duplicate information from the previous section, but if
you're touching many
parts of the code base, its worth explicitly noting the important
changes or how they are tested.
-->

1. adding new geo predicate `vortex_dwithin` in DuckDB, which is
non-throwing and can be pushdown.
2. adding SQL rewrite so for geo native type, `ST_dwithin` is text
rewritten into `vortex_dwithin`.

<!--
Are there any user-facing changes that might require documentation
updates
Is any public API changed?
-->

## Performance

```
┏━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
┃ Query ┃ duckdb:parquet (base) ┃   duckdb:vortex ┃ duckdb:vortex-native ┃
┡━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
│ 1     │               179.6ms │ 132.1ms (0.74x) │       29.0ms (0.16x) │
│ 2     │               268.1ms │ 255.3ms (0.95x) │      312.5ms (1.17x) │
│ 3     │               247.8ms │ 216.4ms (0.87x) │      140.2ms (0.57x) │
│ 4     │               199.6ms │ 146.9ms (0.74x) │      144.5ms (0.72x) │
│ 5     │                 3.40s │   3.43s (1.01x) │        3.03s (0.89x) │
│ 6     │               528.1ms │ 359.2ms (0.68x) │      455.4ms (0.86x) │
│ 7     │               984.0ms │ 983.2ms (1.00x) │      887.7ms (0.90x) │
│ 8     │                 1.08s │ 945.4ms (0.87x) │      997.6ms (0.92x) │
│ 9     │                34.2ms │  33.5ms (0.98x) │       43.6ms (1.27x) │
└───────┴───────────────────────┴─────────────────┴──────────────────────┘
```

Takeaways: Q1 and Q3 is significantly improved due to the single table
geo predicate is pushdown into vortex scan. Q1 is more benefit from the
manually fast path without going through `geo` crate for calculation. Q3
is also potentially can be benefit from manually calculation.

---------

Signed-off-by: Nemo Yu <zyu379@wisc.edu>

Author: HarukiMoriarty

benchmarks/duckdb-bench/src/lib.rs

@@ -78,6 +78,11 @@ impl DuckClient {
         for stmt in &statements {
             self.connection().query(stmt)?;
         }
+        // After `LOAD spatial`, shadow `ST_DWithin` so radius filters push. No-op without it.
+        self.db
+            .as_ref()
+            .vortex_expect("DuckClient database accessed after close")
+            .register_st_dwithin_override()?;
         self.init_sql = statements;
         Ok(())
     }
@@ -127,6 +132,11 @@ impl DuckClient {
                 .vortex_expect("connection just opened")
                 .query(stmt)?;
         }
+        // Re-shadow `ST_DWithin` against the fresh instance.
+        self.db
+            .as_ref()
+            .vortex_expect("database just opened")
+            .register_st_dwithin_override()?;
 
         Ok(())
     }

vortex-duckdb/cpp/expr.cpp

@@ -11,6 +11,19 @@
 #include "duckdb/planner/expression/bound_operator_expression.hpp"
 #include "duckdb/planner/expression/bound_conjunction_expression.hpp"
 
+#include "duckdb/catalog/catalog.hpp"
+#include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp"
+#include "duckdb/common/error_data.hpp"
+#include "duckdb/logging/logger.hpp"
+#include "duckdb/main/capi/capi_internal.hpp"
+#include "duckdb/main/client_context.hpp"
+#include "duckdb/main/connection.hpp"
+#include "duckdb/main/database_manager.hpp"
+#include "duckdb/parser/parsed_data/create_scalar_function_info.hpp"
+#include "duckdb/transaction/meta_transaction.hpp"
+
+#include <exception>
+
 using namespace duckdb;
 
 extern "C" const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func) {
@@ -21,6 +34,52 @@ extern "C" const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func) {
     return func->name.c_str();
 }
 
+extern "C" duckdb_state duckdb_vx_register_st_dwithin_override(duckdb_database ffi_db) {
+    if (!ffi_db) {
+        return DuckDBError;
+    }
+    const DatabaseWrapper &wrapper = *reinterpret_cast<DatabaseWrapper *>(ffi_db);
+    DatabaseInstance &db = *wrapper.database->instance;
+    try {
+        Connection conn(db);
+        ClientContext &context = *conn.context;
+        context.RunFunctionInTransaction([&]() {
+            auto &system = Catalog::GetSystemCatalog(context);
+            auto entry = system.GetEntry<ScalarFunctionCatalogEntry>(context,
+                                                                     DEFAULT_SCHEMA,
+                                                                     "st_dwithin",
+                                                                     OnEntryNotFound::RETURN_NULL);
+            if (!entry) {
+                // No `spatial` loaded, so there is no `ST_DWithin` to override.
+                return;
+            }
+            ScalarFunctionSet set("st_dwithin");
+            for (const auto &overload : entry->functions.functions) {
+                ScalarFunction copy = overload;
+                // Keep the radius as children[2]; spatial's bind folds it into private bind data.
+                copy.bind = nullptr;
+                set.AddFunction(copy);
+            }
+            CreateScalarFunctionInfo info(std::move(set));
+            info.on_conflict = OnCreateConflict::REPLACE_ON_CONFLICT;
+            // `internal` entries are only accepted by the system catalog.
+            info.internal = false;
+            // The user catalog binds ahead of the system catalog, shadowing spatial's entry;
+            // `RestoreStDWithin` rebinds unpushed calls through the original.
+            auto &catalog = Catalog::GetCatalog(context, DatabaseManager::GetDefaultDatabase(context));
+            // Durable catalogs require the modified mark; scalar function entries are never
+            // persisted, so this is metadata-only.
+            MetaTransaction::Get(context).ModifyDatabase(catalog.GetAttached(), DatabaseModificationType());
+            catalog.CreateFunction(context, info);
+        });
+    } catch (const std::exception &e) {
+        ErrorData data(e);
+        DUCKDB_LOG_ERROR(db, "Failed to register the ST_DWithin override:\t" + data.Message());
+        return DuckDBError;
+    }
+    return DuckDBSuccess;
+}
+
 extern "C" const char *duckdb_vx_expr_to_string(duckdb_vx_expr ffi_expr) {
     if (!ffi_expr) {
         return nullptr;

vortex-duckdb/cpp/include/expr.h

@@ -13,6 +13,11 @@ typedef struct duckdb_vx_sfunc_ *duckdb_vx_sfunc;
 
 const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func);
 
+/// Shadow `ST_DWithin` with a copy that keeps the radius as the third argument, so
+/// radius filters can push into Vortex scans. See `RestoreStDWithin` in
+/// scalar_fn_pushdown.hpp for the override/restore example.
+duckdb_state duckdb_vx_register_st_dwithin_override(duckdb_database ffi_db);
+
 typedef struct duckdb_vx_expr_ *duckdb_vx_expr;
 
 /// Return the string representation of the expression. Must be freed with `duckdb_free`.

vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp

@@ -77,6 +77,21 @@ using Projections = unordered_map<TableIndex, LogicalProjection &>;
 
 LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOperatorPtr plan);
 
+/*
+ * We override spatial's `ST_DWithin` with a copy that keeps the radius as a plain third
+ * argument (see expr.h); the original folds it into bind data at bind time. We need the
+ * radius visible to push the filter into a Vortex scan, but spatial's join optimizer only
+ * recognizes the folded 2-argument form. So this pass rebinds join conditions through
+ * spatial's original function and leaves filters alone:
+ *
+ *   FILTER st_dwithin(t.geom, 'POINT(0 0)', 10.0)   -- untouched, pushed into the scan
+ *   JOIN ON st_dwithin(a.geom, b.geom, 10.0)        -- rebound: st_dwithin(a.geom, b.geom)
+ *                                                      with the radius in bind data
+ *
+ * Runs in the pre-optimize hook, before any extension's optimizer pass.
+ */
+void RestoreStDWithin(ClientContext &context, LogicalOperator &plan);
+
 /**
  * Collect fn(col) expressions i.e. expressions where a single function (not
  * a function chain) wraps a single bound column. If "col" is used without

vortex-duckdb/cpp/scalar_fn_pushdown.cpp

@@ -1,6 +1,8 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 #include "duckdb/catalog/catalog.hpp"
+#include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp"
+#include "duckdb/function/function_binder.hpp"
 #include "duckdb/planner/operator/logical_projection.hpp"
 #include "scalar_fn_pushdown.hpp"
 #include "table_function.hpp"
@@ -231,3 +233,67 @@ ScalarFnReplace::ScalarFnReplace(Analyses &analyses, const Projections &projecti
 TableColumnStorageIndex GetAnalysis::StorageIndex(TableColumnScanIndex idx) const {
     return get.GetColumnIds()[idx].GetPrimaryIndex();
 }
+
+namespace {
+
+// See RestoreStDWithin: rebinding through spatial's own entry lets spatial build its own bind
+// data, so nothing here depends on its internals.
+class StDWithinRestore final : public LogicalOperatorVisitor {
+public:
+    explicit StDWithinRestore(ClientContext &context) : context(context) {
+    }
+
+    // Restore join conditions, filters must keep the radius visible so
+    // DuckDB's filter pushdown can offer them to Vortex scans.
+    void VisitOperator(LogicalOperator &op) override {
+        using enum LogicalOperatorType;
+        switch (op.type) {
+        case LOGICAL_COMPARISON_JOIN:
+        case LOGICAL_ANY_JOIN:
+        case LOGICAL_DELIM_JOIN:
+        case LOGICAL_ASOF_JOIN:
+            VisitOperatorExpressions(op);
+            break;
+        default:
+            break;
+        }
+        VisitOperatorChildren(op);
+    }
+
+    ExpressionPtr VisitReplace(BoundFunctionExpression &expr, ExpressionPtr *) override {
+        if (expr.children.size() != 3 || expr.function.name != "st_dwithin") {
+            return nullptr; // Not the override's shape: keep it and descend into its children.
+        }
+        // The system catalog holds spatial's original; the user-catalog override cannot shadow
+        // this lookup.
+        auto original = Catalog::GetSystemCatalog(context).GetEntry<ScalarFunctionCatalogEntry>(
+            context,
+            DEFAULT_SCHEMA,
+            "st_dwithin",
+            OnEntryNotFound::RETURN_NULL);
+        if (!original) {
+            return nullptr;
+        }
+        vector<ExpressionPtr> children;
+        children.reserve(expr.children.size());
+        for (const auto &child : expr.children) {
+            children.push_back(child->Copy());
+        }
+        ErrorData error;
+        FunctionBinder binder(context);
+        auto bound = binder.BindScalarFunction(*original, std::move(children), error);
+        if (!bound) {
+            return nullptr; // No matching overload: keep the executable 3-argument form.
+        }
+        return bound;
+    }
+
+private:
+    ClientContext &context;
+};
+
+} // namespace
+
+void RestoreStDWithin(ClientContext &context, LogicalOperator &plan) {
+    StDWithinRestore(context).VisitOperator(plan);
+}

vortex-duckdb/cpp/vortex_duckdb.cpp

@@ -271,8 +271,13 @@ static void VortexOptimizeFunction(OptimizerExtensionInput &input, unique_ptr<Lo
     plan = TryPushdownScalarFunctions(input.context, std::move(plan));
 }
 
+static void VortexPreOptimizeFunction(OptimizerExtensionInput &input, unique_ptr<LogicalOperator> &plan) {
+    RestoreStDWithin(input.context, *plan);
+}
+
 struct VortexOptimizerExtension final : OptimizerExtension {
-    inline VortexOptimizerExtension() : OptimizerExtension(VortexOptimizeFunction, nullptr, {}) {
+    inline VortexOptimizerExtension()
+        : OptimizerExtension(VortexOptimizeFunction, VortexPreOptimizeFunction, {}) {
     }
 };