Ungrouped aggregate function pushdown for duckdb

vortex-duckdb/build.rs

@@ -27,13 +27,14 @@ const DEFAULT_DUCKDB_VERSION: &str = "1.5.3";
 
 const BUILD_ARTIFACTS: [&str; 3] = ["libduckdb.dylib", "libduckdb.so", "libduckdb_static.a"];
 
-const SOURCE_FILES: [&str; 9] = [
+const SOURCE_FILES: [&str; 10] = [
     "cpp/vortex_duckdb.cpp",
     "cpp/copy_function.cpp",
     "cpp/expr.cpp",
     "cpp/optimizer.cpp",
     "cpp/scalar_fn_pushdown.cpp",
     "cpp/cast_pushdown.cpp",
+    "cpp/aggregate_fn_pushdown.cpp",
     "cpp/table_filter.cpp",
     "cpp/table_function.cpp",
     "cpp/vector.cpp",

vortex-duckdb/cpp/aggregate_fn_pushdown.cpp

@@ -0,0 +1,140 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+#include "aggregate_fn_pushdown.hpp"
+#include "duckdb/planner/expression/bound_aggregate_expression.hpp"
+#include "duckdb/planner/expression/bound_columnref_expression.hpp"
+#include "duckdb/planner/operator/logical_aggregate.hpp"
+#include "optimizer.hpp"
+#include "table_function.hpp"
+
+using enum LogicalOperatorType;
+
+LogicalOperatorPtr TryPushdownAggregateFunctions(ClientContext &context, LogicalOperatorPtr plan) {
+    Analyses analyses;
+    Projections projections;
+    FindGetsAndProjections(*plan, analyses, projections);
+    if (analyses.empty()) {
+        return plan;
+    }
+    return RewriteAggregates(context, std::move(plan), analyses, projections);
+}
+
+LogicalOperatorPtr RewriteAggregates(ClientContext &context,
+                                     LogicalOperatorPtr op,
+                                     Analyses &analyses,
+                                     const Projections &projections) {
+    for (auto &child : op->children) {
+        child = RewriteAggregates(context, std::move(child), analyses, projections);
+    }
+    if (op->type == LOGICAL_AGGREGATE_AND_GROUP_BY) {
+        return TryReplaceAggregate(context, std::move(op), analyses, projections);
+    }
+    return op;
+}
+
+static bool IsUngrouped(const LogicalAggregate &agg) {
+    return agg.groups.empty() && agg.grouping_sets.empty() && agg.grouping_functions.empty() &&
+           !agg.expressions.empty();
+}
+
+constexpr inline idx_t COUNT_STAR_PROJ_IDX = std::numeric_limits<TableColumnStorageIndex>::max();
+
+LogicalOperatorPtr TryReplaceAggregate(ClientContext &context,
+                                       LogicalOperatorPtr op,
+                                       Analyses &analyses,
+                                       const Projections &projections) {
+    LogicalAggregate &agg = op->Cast<LogicalAggregate>();
+    if (!IsUngrouped(agg)) {
+        return op;
+    }
+
+    LogicalGet *const get = GetChildGet(agg);
+    if (get == nullptr) {
+        return op;
+    }
+
+    vector<std::pair<TableColumnScanIndex, const Expression &>> input;
+    const idx_t aggregates_len = agg.expressions.size();
+    input.reserve(aggregates_len);
+
+    for (const auto &expr : agg.expressions) {
+        if (expr->GetExpressionClass() != ExpressionClass::BOUND_AGGREGATE) {
+            return op;
+        }
+        const auto &bound_aggr = expr->Cast<BoundAggregateExpression>();
+        if (bound_aggr.IsDistinct() || bound_aggr.filter != nullptr || bound_aggr.order_bys != nullptr) {
+            return op;
+        }
+
+        if (bound_aggr.function.name == "count_star") {
+            input.emplace_back(COUNT_STAR_PROJ_IDX, *expr);
+            continue;
+        }
+
+        if (bound_aggr.children.size() != 1 ||
+            bound_aggr.children[0]->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) {
+            return op;
+        }
+        const auto &bound_col = bound_aggr.children[0]->Cast<BoundColumnRefExpression>();
+        const auto binding = Resolve(bound_col.binding, analyses, projections);
+        if (!binding || &binding->analysis.get != get) {
+            return op;
+        }
+        input.emplace_back(binding->column_index, *expr);
+    }
+
+    if (!aggregate_pushdown(context, {*get, input})) {
+        return op;
+    }
+
+    // GET now returns one column per aggregate. Expand existing columns
+    auto &column_ids = get->GetMutableColumnIds();
+    get->types.resize(aggregates_len);
+    get->returned_types.resize(aggregates_len);
+    column_ids.resize(aggregates_len);
+
+    vector<string> names(aggregates_len); // need a copy because we reference original names
+
+    for (idx_t i = 0; i < aggregates_len; i++) {
+        const auto &[column_index, expr] = input[i];
+        if (column_index == COUNT_STAR_PROJ_IDX) {
+            names[i] = "count_star()";
+        } else {
+            const TableColumnStorageIndex storage_index = get->GetColumnIds()[column_index].GetPrimaryIndex();
+            names[i] = get->names[storage_index];
+        }
+        get->types[i] = expr.return_type;
+        get->returned_types[i] = expr.return_type;
+        column_ids[i] = ColumnIndex {i};
+    }
+    get->names = std::move(names);
+    get->projection_ids.clear();
+    get->table_index = agg.aggregate_index;
+
+    unique_ptr<LogicalOperator> &child = agg.children[0];
+    if (child->type == LOGICAL_GET) {
+        return std::move(child);
+    }
+    D_ASSERT(child->type == LOGICAL_PROJECTION);
+    D_ASSERT(child->children.size() == 1);
+    D_ASSERT(child->children[0]->type == LOGICAL_GET);
+    return std::move(child->children[0]);
+}
+
+LogicalGet *GetChildGet(const LogicalAggregate &agg) {
+    if (agg.children.size() != 1) {
+        return nullptr;
+    }
+    LogicalOperator &child = *agg.children[0];
+    LogicalOperator *op;
+    if (child.type == LOGICAL_GET) {
+        op = &child;
+    } else if (child.type == LOGICAL_PROJECTION && child.children.size() == 1 &&
+               child.children[0]->type == LOGICAL_GET) {
+        op = child.children[0].get();
+    } else {
+        return nullptr;
+    }
+    LogicalGet &get = op->Cast<LogicalGet>();
+    return get.function.bind == duckdb_vx_table_function_bind ? &get : nullptr;
+}

vortex-duckdb/cpp/expr.cpp

@@ -3,11 +3,13 @@
 
 #include "expr.h"
 #include "duckdb/function/scalar_function.hpp"
+#include "duckdb/function/aggregate_function.hpp"
 #include "duckdb/planner/expression/bound_between_expression.hpp"
 #include "duckdb/planner/expression/bound_cast_expression.hpp"
 #include "duckdb/planner/expression/bound_columnref_expression.hpp"
 #include "duckdb/planner/expression/bound_comparison_expression.hpp"
 #include "duckdb/planner/expression/bound_constant_expression.hpp"
+#include "duckdb/planner/expression/bound_aggregate_expression.hpp"
 #include "duckdb/planner/expression/bound_function_expression.hpp"
 #include "duckdb/planner/expression/bound_operator_expression.hpp"
 #include "duckdb/planner/expression/bound_conjunction_expression.hpp"
@@ -81,6 +83,11 @@ extern "C" duckdb_state duckdb_vx_register_st_dwithin_override(duckdb_database f
     return DuckDBSuccess;
 }
 
+extern "C" const char *duckdb_vx_agg_func_name(duckdb_vx_agg_func ffi) {
+    D_ASSERT(ffi);
+    return reinterpret_cast<AggregateFunction *>(ffi)->name.c_str();
+}
+
 extern "C" const char *duckdb_vx_expr_to_string(duckdb_vx_expr ffi_expr) {
     if (!ffi_expr) {
         return nullptr;
@@ -92,7 +99,6 @@ extern "C" const char *duckdb_vx_expr_to_string(duckdb_vx_expr ffi_expr) {
     return result;
 }
 
-//! Create a DuckDB vortex error.
 extern "C" void duckdb_vx_destroy_expr(duckdb_vx_expr *ffi_expr) {
     auto expr = reinterpret_cast<Expression *>(ffi_expr);
     delete expr;
@@ -201,3 +207,9 @@ extern "C" bool duckdb_vx_expr_get_bound_cast_is_try(duckdb_vx_expr ffi_expr) {
     auto &expr = reinterpret_cast<Expression *>(ffi_expr)->Cast<BoundCastExpression>();
     return expr.try_cast;
 }
+
+extern "C" duckdb_vx_agg_func duckdb_vx_expr_get_bound_aggregate_function(duckdb_vx_expr ffi_expr) {
+    D_ASSERT(ffi_expr);
+    auto &expr = reinterpret_cast<Expression *>(ffi_expr)->Cast<BoundAggregateExpression>();
+    return reinterpret_cast<duckdb_vx_agg_func>(&expr.function);
+}

vortex-duckdb/cpp/include/aggregate_fn_pushdown.hpp

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+#pragma once
+#include "optimizer.hpp"
+#include "duckdb/optimizer/optimizer_extension.hpp"
+
+using namespace duckdb;
+
+// Push UNGROUPED_AGGREGATE's of form agg(T) and count_star() into GET.
+LogicalOperatorPtr TryPushdownAggregateFunctions(ClientContext &context, LogicalOperatorPtr plan);
+
+LogicalOperatorPtr RewriteAggregates(ClientContext &context,
+                                     LogicalOperatorPtr op,
+                                     Analyses &analyses,
+                                     const Projections &projections);
+
+LogicalOperatorPtr TryReplaceAggregate(ClientContext &context,
+                                       LogicalOperatorPtr op,
+                                       Analyses &analyses,
+                                       const Projections &projections);
+
+// return GET for UNGROUPED_AGGREGATE -> [GET] or for UNGROUPED_AGGREGATE ->
+// PROJECTION -> [GET], nullptr if not found.
+LogicalGet *GetChildGet(const LogicalAggregate &agg);

vortex-duckdb/cpp/include/expr.h

@@ -10,6 +10,7 @@ extern "C" {
 #endif
 
 typedef struct duckdb_vx_sfunc_ *duckdb_vx_sfunc;
+typedef struct duckdb_vx_agg_func_ *duckdb_vx_agg_func;
 
 const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func);
 
@@ -20,6 +21,8 @@ duckdb_state duckdb_vx_register_st_dwithin_override(duckdb_database ffi_db);
 
 typedef struct duckdb_vx_expr_ *duckdb_vx_expr;
 
+const char *duckdb_vx_agg_func_name(duckdb_vx_agg_func func);
+
 /// Return the string representation of the expression. Must be freed with `duckdb_free`.
 const char *duckdb_vx_expr_to_string(duckdb_vx_expr expr);
 
@@ -273,6 +276,8 @@ duckdb_vx_expr duckdb_vx_expr_get_bound_cast_child(duckdb_vx_expr expr);
 
 bool duckdb_vx_expr_get_bound_cast_is_try(duckdb_vx_expr expr);
 
+duckdb_vx_agg_func duckdb_vx_expr_get_bound_aggregate_function(duckdb_vx_expr expr);
+
 #ifdef __cplusplus /* End C ABI */
 }
 #endif

vortex-duckdb/cpp/include/table_function.h

@@ -88,6 +88,10 @@ typedef struct {
 
 duckdb_state duckdb_vx_register_table_functions(duckdb_database ffi_db);
 
+typedef struct duckdb_vx_agg_input_ *duckdb_vx_agg_input;
+idx_t duckdb_vx_aggregate_len(duckdb_vx_agg_input ffi);
+duckdb_vx_expr duckdb_vx_aggregate_at(duckdb_vx_agg_input ffi, idx_t index, idx_t *proj_idx);
+
 #ifdef __cplusplus
 }
 #endif

vortex-duckdb/cpp/include/table_function.hpp

@@ -25,3 +25,11 @@ struct TableFunctionProjectionExpressionInput {
 // true if we can push down the expression, false otherwise
 bool projection_expression_pushdown(duckdb::ClientContext &context,
                                     const TableFunctionProjectionExpressionInput &input);
+
+struct TableFunctionUngroupedAggregateInput {
+    const duckdb::LogicalGet &get;
+    // Column scan index -> aggregate expression
+    const duckdb::vector<std::pair<idx_t, const duckdb::Expression &>> &projections;
+};
+
+bool aggregate_pushdown(duckdb::ClientContext &context, const TableFunctionUngroupedAggregateInput &input);

vortex-duckdb/cpp/table_function.cpp

@@ -4,6 +4,7 @@
 #include "data.hpp"
 #include "error.hpp"
 #include "table_function.hpp"
+#include "expr.h"
 #include "vortex_duckdb.h"
 #include "table_function.h"
 #include "vortex.h"
@@ -171,12 +172,16 @@ struct CTableBindResult {
     vector<string> &names;
 };
 
+// This is a flaw of Duckdb API which doesn't allow passing non-const
+// expressions. We never modify the value on Rust side.
+static duckdb_vx_expr get_ffi_expr(const Expression &expr) {
+    return reinterpret_cast<duckdb_vx_expr>(const_cast<Expression *>(&expr));
+}
+
 bool projection_expression_pushdown(ClientContext &, const TableFunctionProjectionExpressionInput &input) {
     const auto &bind = input.get.bind_data->Cast<CTableBindData>();
 
-    // This is a flaw of Duckdb API which doesn't allow passing non-const
-    // expressions. We never modify the value on Rust side.
-    auto ffi_expr = reinterpret_cast<duckdb_vx_expr>(const_cast<Expression *>(&input.expression));
+    duckdb_vx_expr ffi_expr = get_ffi_expr(input.expression);
     void *const ffi_bind = bind.ffi_data->DataPtr();
     duckdb_vx_error error_out = nullptr;
 
@@ -191,6 +196,33 @@ bool projection_expression_pushdown(ClientContext &, const TableFunctionProjecti
     return ret;
 }
 
+extern "C" {
+idx_t duckdb_vx_aggregate_len(duckdb_vx_agg_input ffi_input) {
+    return reinterpret_cast<const TableFunctionUngroupedAggregateInput *>(ffi_input)->projections.size();
+}
+
+duckdb_vx_expr duckdb_vx_aggregate_at(duckdb_vx_agg_input ffi_input, idx_t i, idx_t *proj_idx) {
+    const auto &input = *reinterpret_cast<const TableFunctionUngroupedAggregateInput *>(ffi_input);
+    const auto &[scan_index, expr] = input.projections[i];
+    *proj_idx = scan_index == COUNT_STAR_PROJ_IDX ? scan_index
+                                                  : input.get.GetColumnIds()[scan_index].GetPrimaryIndex();
+    return get_ffi_expr(expr);
+}
+}
+
+bool aggregate_pushdown(ClientContext &, const TableFunctionUngroupedAggregateInput &input) {
+    const auto &bind = input.get.bind_data->Cast<CTableBindData>();
+    void *const ffi_bind = bind.ffi_data->DataPtr();
+    duckdb_vx_error error_out = nullptr;
+    const auto ffi_input =
+        reinterpret_cast<duckdb_vx_agg_input>(const_cast<TableFunctionUngroupedAggregateInput *>(&input));
+    const bool res = duckdb_table_function_pushdown_projection_aggregates(ffi_bind, ffi_input, &error_out);
+    if (error_out) {
+        throw BinderException(IntoErrString(error_out));
+    }
+    return res;
+}
+
 /**
  * Called for every new query. For example, if there is a VIEW over *.vortex,
  * and after a query another file is added matching the glob, for second query
@@ -238,10 +270,11 @@ unique_ptr<GlobalTableFunctionState> c_init_global(ClientContext &context, Table
 }
 
 unique_ptr<LocalTableFunctionState>
-init_local(ExecutionContext &, TableFunctionInitInput &, GlobalTableFunctionState *global_state) {
+init_local(ExecutionContext &, TableFunctionInitInput &input, GlobalTableFunctionState *global_state) {
+    const void *const ffi_bind = input.bind_data->Cast<CTableBindData>().ffi_data->DataPtr();
     void *const ffi_global = global_state->Cast<CTableGlobalData>().ffi_data->DataPtr();
 
-    duckdb_vx_data ffi_local_data = duckdb_table_function_init_local(ffi_global);
+    duckdb_vx_data ffi_local_data = duckdb_table_function_init_local(ffi_bind, ffi_global);
     auto cdata = unique_ptr<CData>(reinterpret_cast<CData *>(ffi_local_data));
     return make_uniq<CTableLocalData>(std::move(cdata));
 }

vortex-duckdb/cpp/vortex_duckdb.cpp

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+#include "aggregate_fn_pushdown.hpp"
 #include "data.hpp"
 #include "error.hpp"
 #include "scalar_fn_pushdown.hpp"
@@ -271,6 +272,7 @@ extern "C" duckdb_blob duckdb_vx_value_get_geometry(duckdb_value value) {
 static void VortexOptimizeFunction(OptimizerExtensionInput &input, unique_ptr<LogicalOperator> &plan) {
     plan = TryPushdown<CastCollect, CastReplace>(input.context, std::move(plan));
     plan = TryPushdown<ScalarFnCollect, ScalarFnReplace>(input.context, std::move(plan));
+    plan = TryPushdownAggregateFunctions(input.context, std::move(plan));
 }
 
 static void VortexPreOptimizeFunction(OptimizerExtensionInput &input, unique_ptr<LogicalOperator> &plan) {

vortex-duckdb/include/vortex.h

@@ -11,6 +11,8 @@
 
 #pragma once
 
+#define COUNT_STAR_PROJ_IDX UINT64_MAX
+
 #ifdef __cplusplus
 extern "C" {
 #endif // __cplusplus
@@ -40,6 +42,11 @@ bool duckdb_table_function_pushdown_projection_expression(void *bind_data,
                                                           size_t column_id,
                                                           duckdb_vx_error *error_out);
 
+extern
+bool duckdb_table_function_pushdown_projection_aggregates(void *bind_data,
+                                                          duckdb_vx_agg_input input,
+                                                          duckdb_vx_error *error_out);
+
 extern
 void duckdb_table_function_scan(void *global_init_data,
                                 void *local_init_data,
@@ -56,7 +63,9 @@ extern
 duckdb_vx_data duckdb_table_function_init_global(const duckdb_vx_tfunc_init_input *init_input,
                                                  duckdb_vx_error *error_out);
 
-extern duckdb_vx_data duckdb_table_function_init_local(void *global_init_data);
+extern
+duckdb_vx_data duckdb_table_function_init_local(const void *bind_data,
+                                                void *global_init_data);
 
 extern
 duckdb_vx_data duckdb_table_function_bind(duckdb_vx_tfunc_bind_input bind_input,