Ungrouped aggregate function pushdown for duckdb

vortex-duckdb/build.rs

@@ -27,11 +27,12 @@ const DEFAULT_DUCKDB_VERSION: &str = "1.5.3";
 
 const BUILD_ARTIFACTS: [&str; 3] = ["libduckdb.dylib", "libduckdb.so", "libduckdb_static.a"];
 
-const SOURCE_FILES: [&str; 7] = [
+const SOURCE_FILES: [&str; 8] = [
     "cpp/vortex_duckdb.cpp",
     "cpp/copy_function.cpp",
     "cpp/expr.cpp",
     "cpp/scalar_fn_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,136 @@
+// 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/operator/logical_aggregate.hpp"
+#include "scalar_fn_pushdown.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<TableColumnStorageIndex, const Expression &>> input;
+    const idx_t N = agg.expressions.size();
+    input.reserve(N);
+
+    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;
+        }
+        const TableColumnStorageIndex storage_index = binding->analysis.StorageIndex(binding->column_index);
+        input.emplace_back(storage_index, *expr);
+    }
+
+    if (!aggregate_pushdown(context, {*get, input})) {
+        return op;
+    }
+
+    // GET now returns one column per aggregate. Expand column into multiple
+    // column if there are many aggregates per column.
+    auto &column_ids = get->GetMutableColumnIds();
+    get->types.resize(N);
+    get->returned_types.resize(N);
+    column_ids.resize(N);
+
+    vector<string> names(N); // need a copy because we reference original names
+
+    for (idx_t i = 0; i < N; i++) {
+        const auto &[storage_index, expr] = input[i];
+        names[i] = storage_index == COUNT_STAR_PROJ_IDX ? "count_star()" : 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,10 +3,12 @@
 
 #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_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"
@@ -21,6 +23,11 @@ extern "C" const char *duckdb_vx_sfunc_name(duckdb_vx_sfunc ffi_func) {
     return func->name.c_str();
 }
 
+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;
@@ -129,3 +136,9 @@ extern "C" void duckdb_vx_expr_get_bound_function(duckdb_vx_expr ffi_expr,
     out->scalar_function = reinterpret_cast<duckdb_vx_sfunc>(&expr.function);
     out->bind_info = expr.bind_info.get();
 }
+
+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 "scalar_fn_pushdown.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,11 +10,14 @@ 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);
 
 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);
 
@@ -264,6 +267,8 @@ typedef struct {
 
 void duckdb_vx_expr_get_bound_function(duckdb_vx_expr expr, duckdb_vx_expr_bound_function *out);
 
+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_i(duckdb_vx_agg_input ffi, idx_t i, idx_t *proj_idx);
+
 #ifdef __cplusplus
 }
 #endif

vortex-duckdb/cpp/include/table_function.hpp

@@ -25,3 +25,12 @@ struct TableFunctionProjectionExpressionInput {
 // true if we can push down the expression, false otherwise
 bool projection_expression_pushdown(duckdb::ClientContext &context,
                                     const TableFunctionProjectionExpressionInput &input);
+
+// API is subject to change, but you can push only none or all aggregates since
+// this changes output chunk cardinality
+struct TableFunctionUngroupedAggregateInput {
+    const duckdb::LogicalGet &get;
+    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;
 }
 
+using Projections = vector<std::pair<idx_t, const Expression &>>;
+
+extern "C" {
+idx_t duckdb_vx_aggregate_len(duckdb_vx_agg_input ffi) {
+    return reinterpret_cast<const Projections *>(ffi)->size();
+}
+
+duckdb_vx_expr duckdb_vx_aggregate_i(duckdb_vx_agg_input ffi, idx_t i, idx_t *proj_idx) {
+    const Projections &projections = *reinterpret_cast<const Projections *>(ffi);
+    *proj_idx = projections[i].first;
+    return get_ffi_expr(projections[i].second);
+}
+}
+
+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<Projections *>(&input.projections));
+    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"
@@ -269,6 +270,7 @@ extern "C" duckdb_blob duckdb_vx_value_get_geometry(duckdb_value value) {
 
 static void VortexOptimizeFunction(OptimizerExtensionInput &input, unique_ptr<LogicalOperator> &plan) {
     plan = TryPushdownScalarFunctions(input.context, std::move(plan));
+    plan = TryPushdownAggregateFunctions(input.context, std::move(plan));
 }
 
 struct VortexOptimizerExtension final : OptimizerExtension {

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,