move optimizer code to scalar_fn_pushdown (#8605)

This is preparatory work for aggregate pushdown

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

vortex-array/src/executor.rs

@@ -368,7 +368,7 @@ impl ExecutionCtx {
     /// Use the [`format_args!`] macro to create the `msg` argument.
     pub fn log(&mut self, msg: fmt::Arguments<'_>) {
         #[cfg(debug_assertions)]
-        if tracing::enabled!(tracing::Level::DEBUG) {
+        if tracing::enabled!(tracing::Level::TRACE) {
             let formatted = format!(" - {msg}");
             tracing::trace!("exec[{}]: {formatted}", self.id);
             self.ops.push(formatted);

vortex-duckdb/build.rs

@@ -31,7 +31,7 @@ const SOURCE_FILES: [&str; 7] = [
     "cpp/vortex_duckdb.cpp",
     "cpp/copy_function.cpp",
     "cpp/expr.cpp",
-    "cpp/optimizer.cpp",
+    "cpp/scalar_fn_pushdown.cpp",
     "cpp/table_filter.cpp",
     "cpp/table_function.cpp",
     "cpp/vector.cpp",
@@ -175,14 +175,13 @@ const DUCKDB_C_API_FUNCTIONS: [&str; 133] = [
     "duckdb_vector_size",
 ];
 
-const DUCKDB_C_API_HEADERS: [&str; 7] = [
+const DUCKDB_C_API_HEADERS: [&str; 6] = [
     "cpp/include/vortex_duckdb.h",
     "cpp/include/expr.h",
     "cpp/include/table_filter.h",
     "cpp/include/vector.h",
     "cpp/include/copy_function.h",
     "cpp/include/table_function.h",
-    "cpp/include/optimizer.h",
 ];
 
 const DOWNLOAD_MAX_RETRIES: i32 = 3;

vortex-duckdb/cpp/include/optimizer.h

@@ -1,15 +1,14 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
-
 #pragma once
+#include "duckdb.h"
 
 #include "duckdb/optimizer/optimizer_extension.hpp"
 #include "duckdb/planner/expression/bound_columnref_expression.hpp"
 #include "duckdb/planner/expression/bound_function_expression.hpp"
 #include "duckdb/planner/operator/logical_get.hpp"
 #include <optional>
 
-// Only one consumer of this header file, so "using" is fine
 using namespace duckdb;
 
 using ExpressionPtr = unique_ptr<Expression>;
@@ -50,6 +49,8 @@ struct GetAnalysis {
      * or without function application in the query plan.
      */
     unordered_map<TableColumnScanIndex, const BoundFunctionExpression *> col_to_fn;
+
+    TableColumnStorageIndex StorageIndex(TableColumnScanIndex idx) const;
 };
 
 using Analyses = unordered_map<TableIndex, GetAnalysis>;
@@ -74,6 +75,8 @@ using Analyses = unordered_map<TableIndex, GetAnalysis>;
  */
 using Projections = unordered_map<TableIndex, LogicalProjection &>;
 
+LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOperatorPtr 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
@@ -105,14 +108,6 @@ struct ScalarFnReplace final : LogicalOperatorVisitor {
 
 void FindGetsAndProjections(LogicalOperator &op, Analyses &analyses, Projections &aliases);
 
-LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOperatorPtr plan);
-void VortexOptimizeFunction(OptimizerExtensionInput &input, LogicalOperatorPtr &plan);
-
-struct VortexOptimizerExtension final : OptimizerExtension {
-    inline VortexOptimizerExtension() : OptimizerExtension(VortexOptimizeFunction, nullptr, {}) {
-    }
-};
-
 struct GetBinding {
     GetAnalysis &analysis;
     TableColumnScanIndex column_index;

vortex-duckdb/cpp/include/optimizer.hpp …uckdb/cpp/include/scalar_fn_pushdown.hppvortex-duckdb/cpp/include/optimizer.hpp renamed to vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp vortex-duckdb/cpp/include/optimizer.hpp renamed to vortex-duckdb/cpp/include/scalar_fn_pushdown.hpp

@@ -43,6 +43,8 @@ duckdb_logical_type duckdb_vx_create_geometry(const char *crs);
 
 duckdb_state duckdb_vx_register_scan_replacement(duckdb_database duckdb_database);
 
+duckdb_state duckdb_vx_optimizer_extension_register(duckdb_database ffi_db);
+
 /// Creates a new reusable dictionary from a logical type and size.
 duckdb_vx_reusable_dict duckdb_vx_reusable_dict_create(duckdb_logical_type logical_type, idx_t size);
 

vortex-duckdb/cpp/include/vortex_duckdb.h

@@ -1,11 +1,9 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
-#include "optimizer.hpp"
-#include "table_function.hpp"
 #include "duckdb/catalog/catalog.hpp"
-#include "duckdb/main/config.hpp"
-#include "duckdb/main/capi/capi_internal.hpp"
 #include "duckdb/planner/operator/logical_projection.hpp"
+#include "scalar_fn_pushdown.hpp"
+#include "table_function.hpp"
 #include <optional>
 
 /**
@@ -15,24 +13,6 @@
  * may produce conflicts (e.g. WHERE prefix("str", 'h')).
  */
 
-extern "C" duckdb_state duckdb_vx_optimizer_extension_register(duckdb_database ffi_db) {
-    D_ASSERT(ffi_db);
-    const DatabaseWrapper &wrapper = *reinterpret_cast<DatabaseWrapper *>(ffi_db);
-    DatabaseInstance &db = *wrapper.database->instance;
-    try {
-        DBConfig::GetConfig(db).GetCallbackManager().Register(VortexOptimizerExtension());
-    } catch (const std::exception &e) {
-        ErrorData data(e);
-        DUCKDB_LOG_ERROR(db, "Failed to create Vortex optimizer extension:\t" + data.Message());
-        return DuckDBError;
-    }
-    return DuckDBSuccess;
-}
-
-void VortexOptimizeFunction(OptimizerExtensionInput &input, LogicalOperatorPtr &plan) {
-    plan = TryPushdownScalarFunctions(input.context, std::move(plan));
-}
-
 LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOperatorPtr plan) {
     Analyses analyses;
     Projections projections;
@@ -48,8 +28,7 @@ LogicalOperatorPtr TryPushdownScalarFunctions(ClientContext &context, LogicalOpe
             if (expr == nullptr) { // Conflict for column
                 continue;
             }
-            const TableColumnStorageIndex storage_index =
-                analysis.get.GetColumnIds()[column_index].GetPrimaryIndex();
+            const TableColumnStorageIndex storage_index = analysis.StorageIndex(column_index);
             TableFunctionProjectionExpressionInput input {analysis.get, *expr, storage_index};
             if (projection_expression_pushdown(context, input)) {
                 analysis.get.types[column_index] = expr->return_type;
@@ -248,3 +227,7 @@ ScalarFnCollect::ScalarFnCollect(Analyses &analyses, const Projections &projecti
 ScalarFnReplace::ScalarFnReplace(Analyses &analyses, const Projections &projections)
     : analyses(analyses), projections(projections) {
 }
+
+TableColumnStorageIndex GetAnalysis::StorageIndex(TableColumnScanIndex idx) const {
+    return get.GetColumnIds()[idx].GetPrimaryIndex();
+}

vortex-duckdb/cpp/optimizer.cpp vortex-duckdb/cpp/scalar_fn_pushdown.cppvortex-duckdb/cpp/optimizer.cpp renamed to vortex-duckdb/cpp/scalar_fn_pushdown.cpp

@@ -3,8 +3,10 @@
 
 #include "data.hpp"
 #include "error.hpp"
+#include "scalar_fn_pushdown.hpp"
 #include "vortex_duckdb.h"
 
+#include "duckdb/catalog/catalog.hpp"
 #include "duckdb/common/assert.hpp"
 #include "duckdb/common/types.hpp"
 #include "duckdb/common/types/data_chunk.hpp"
@@ -15,6 +17,7 @@
 #include "duckdb/main/client_context.hpp"
 #include "duckdb/main/config.hpp"
 #include "duckdb/main/connection.hpp"
+#include "duckdb/optimizer/optimizer_extension.hpp"
 #include "duckdb/parser/expression/constant_expression.hpp"
 #include "duckdb/parser/expression/function_expression.hpp"
 #include "duckdb/parser/tableref/table_function_ref.hpp"
@@ -263,3 +266,26 @@ extern "C" duckdb_blob duckdb_vx_value_get_geometry(duckdb_value value) {
     }
     return {buf, size};
 }
+
+static void VortexOptimizeFunction(OptimizerExtensionInput &input, unique_ptr<LogicalOperator> &plan) {
+    plan = TryPushdownScalarFunctions(input.context, std::move(plan));
+}
+
+struct VortexOptimizerExtension final : OptimizerExtension {
+    inline VortexOptimizerExtension() : OptimizerExtension(VortexOptimizeFunction, nullptr, {}) {
+    }
+};
+
+extern "C" duckdb_state duckdb_vx_optimizer_extension_register(duckdb_database ffi_db) {
+    D_ASSERT(ffi_db);
+    const DatabaseWrapper &wrapper = *reinterpret_cast<DatabaseWrapper *>(ffi_db);
+    DatabaseInstance &db = *wrapper.database->instance;
+    try {
+        DBConfig::GetConfig(db).GetCallbackManager().Register(VortexOptimizerExtension());
+    } catch (const std::exception &e) {
+        ErrorData data(e);
+        DUCKDB_LOG_ERROR(db, "Failed to create Vortex optimizer extension:\t" + data.Message());
+        return DuckDBError;
+    }
+    return DuckDBSuccess;
+}