Cast pushdown for duckdb

vortex-duckdb/build.rs

@@ -27,11 +27,13 @@ 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; 9] = [
     "cpp/vortex_duckdb.cpp",
     "cpp/copy_function.cpp",
     "cpp/expr.cpp",
+    "cpp/optimizer.cpp",
     "cpp/scalar_fn_pushdown.cpp",
+    "cpp/cast_pushdown.cpp",
     "cpp/table_filter.cpp",
     "cpp/table_function.cpp",
     "cpp/vector.cpp",

vortex-duckdb/cpp/cast_pushdown.cpp

@@ -0,0 +1,171 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+#include "cast_pushdown.hpp"
+#include "table_function.hpp"
+
+#include "duckdb/planner/operator/logical_get.hpp"
+#include "duckdb/planner/operator/logical_projection.hpp"
+#include "duckdb/planner/expression/bound_cast_expression.hpp"
+#include "duckdb/planner/expression/bound_columnref_expression.hpp"
+
+// A GET reachable through a single-child chain of filters/projections. A join
+// (or any other multi-child operator) breaks the chain.
+// See test/sql/copy/csv/test_insert_into_types.test in duckdb (cast not pushed past a join)
+static bool ReachesPushdownGet(const LogicalOperator &op) {
+    const LogicalOperator *cur = &op;
+    while (cur->children.size() == 1) {
+        cur = cur->children[0].get();
+        switch (cur->type) {
+        case LogicalOperatorType::LOGICAL_GET:
+            return cur->Cast<LogicalGet>().function.bind == duckdb_vx_table_function_bind;
+        case LogicalOperatorType::LOGICAL_FILTER:
+        case LogicalOperatorType::LOGICAL_PROJECTION:
+            continue;
+        default:
+            return false;
+        }
+    }
+    return false;
+}
+
+void CastCollect::VisitOperator(LogicalOperator &op) {
+    /*
+     * Logical projection expressions are columns which reference underlying
+     * GETs. Don't process them, as they would add conflicts for every column
+     * used in projection. Example: PROJECTION(col) -> GET(col). We don't want
+     * to visit BoundColumnRefExpression in PROJECTION to avoid registering a
+     * non-existent conflict.
+     *
+     * However, CastReplace will visit them because we need to update their
+     * types if pushdown succeeded.
+     */
+    if (op.type != LogicalOperatorType::LOGICAL_PROJECTION) {
+        return LogicalOperatorVisitor::VisitOperator(op);
+    }
+    auto &projection = op.Cast<LogicalProjection>();
+
+    // Only push casts from a projection that forwards just column refs and
+    // casts and reaches a GET without a join in between. A constant or other
+    // expression makes the projection ineligible.
+    // See test/sql/copy/csv/test_csv_error_message_type.test (top-level cast
+    // to VARCHAR must still push) and test_large_integer_detection.test (a
+    // nested cast to VARCHAR must not) in duckdb.
+    bool clean = ReachesPushdownGet(projection);
+    for (const auto &e : projection.expressions) {
+        switch (e->GetExpressionClass()) {
+        case ExpressionClass::BOUND_COLUMN_REF:
+        case ExpressionClass::BOUND_CAST:
+            continue;
+        default:
+            clean = false;
+            break;
+        }
+    }
+    if (clean) {
+        for (const auto &e : projection.expressions) {
+            if (e->GetExpressionClass() == ExpressionClass::BOUND_CAST) {
+                top_level_casts.insert(e.get());
+            }
+        }
+    }
+    if (projections.count(projection.table_index)) {
+        VisitOperatorChildren(op);
+        return;
+    }
+
+    LogicalOperatorVisitor::VisitOperator(op);
+}
+
+ExpressionPtr CastCollect::VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) {
+    if (const auto binding = Resolve(expr.binding, analyses, projections)) {
+        // Column is used without cast applied to it, register a conflict.
+        // Not emplace() as we need to update the value if it was present
+        binding->analysis.col_to_expr[binding->column_index] = nullptr;
+    }
+    return std::move(*ptr);
+}
+
+ExpressionPtr CastCollect::VisitReplace(BoundCastExpression &expr, ExpressionPtr *ptr) {
+    if (expr.child->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) {
+        // Descend into children so e.g. fn(col, other) still sees "col" and
+        // registers a conflict
+        return nullptr;
+    }
+    const auto &bound_col = expr.child->Cast<BoundColumnRefExpression>();
+    const auto binding = Resolve(bound_col.binding, analyses, projections);
+    if (!binding) {
+        return nullptr;
+    }
+    auto &col_to_expr = binding->analysis.col_to_expr;
+
+    if (auto it = col_to_expr.find(binding->column_index); it == col_to_expr.end()) {
+        // Only a top-level projection cast starts a candidate.
+        if (top_level_casts.count(&expr)) {
+            col_to_expr.emplace(binding->column_index, &expr);
+        }
+    } else if (it->second == nullptr ||
+               it->second->Cast<BoundCastExpression>().return_type != expr.return_type ||
+               // TODO(myrrc) this line needs upstreaming
+               it->second->Cast<BoundCastExpression>().try_cast != expr.try_cast) {
+        // Different target type, or already a conflict.
+        it->second = nullptr;
+    }
+
+    return std::move(*ptr);
+}
+
+ExpressionPtr CastReplace::VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) {
+    const auto binding = Resolve(expr.binding, analyses, projections);
+    if (!binding) {
+        return std::move(*ptr);
+    }
+
+    const auto &[analysis, column_index, projection] = *binding;
+    if (CanPushdownColumn(analysis, column_index)) {
+        const idx_t storage_index = analysis.get.GetColumnIds()[column_index].GetPrimaryIndex();
+        const LogicalType return_type = analysis.get.returned_types[storage_index];
+        expr.return_type = return_type;
+        // LogicalProjection types are resolved by calling
+        // LogicalProjection::ResolveTypes, so we need to check whether types in
+        // projection have been resolved, and updated them only if needed.
+        if (projection != nullptr && !projection->types.empty()) {
+            projection->types[column_index] = return_type;
+        }
+    }
+
+    return std::move(*ptr);
+}
+
+ExpressionPtr CastReplace::VisitReplace(BoundCastExpression &expr, ExpressionPtr *ptr) {
+    if (expr.child->GetExpressionType() != ExpressionType::BOUND_COLUMN_REF) {
+        return nullptr; // Same as in ScalarFnCollect::VisitReplace
+    }
+    auto &bound_col_base = expr.child;
+    const auto &bound_col = bound_col_base->Cast<BoundColumnRefExpression>();
+    const auto binding = Resolve(bound_col.binding, analyses, projections);
+    if (!binding) {
+        return nullptr;
+    }
+
+    const auto &[analysis, column_index, projection] = *binding;
+    if (!CanPushdownColumn(analysis, column_index)) {
+        return std::move(*ptr);
+    }
+
+    const idx_t storage_index = analysis.get.GetColumnIds()[column_index].GetPrimaryIndex();
+    const LogicalType return_type = analysis.get.returned_types[storage_index];
+    bound_col_base->return_type = return_type;
+    // Same as in CastReplace::VisitReplace(BoundColumnRefExpression)
+    if (projection != nullptr && !projection->types.empty()) {
+        projection->types[column_index] = return_type;
+    }
+    return std::move(bound_col_base);
+}
+
+CastCollect::CastCollect(Analyses &analyses, const Projections &projections)
+    : analyses(analyses), projections(projections) {
+}
+
+CastReplace::CastReplace(Analyses &analyses, const Projections &projections)
+    : analyses(analyses), projections(projections) {
+}

vortex-duckdb/cpp/expr.cpp

@@ -2,8 +2,10 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 #include "expr.h"
+#include "duckdb/common/type_visitor.hpp"
 #include "duckdb/function/scalar_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"
@@ -129,3 +131,23 @@ 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_expr duckdb_vx_expr_get_bound_cast_child(duckdb_vx_expr ffi_expr) {
+    D_ASSERT(ffi_expr);
+    auto &expr = reinterpret_cast<Expression *>(ffi_expr)->Cast<BoundCastExpression>();
+    return reinterpret_cast<duckdb_vx_expr>(expr.child.get());
+}
+
+extern "C" bool duckdb_vx_expr_get_bound_cast_is_try(duckdb_vx_expr ffi_expr) {
+    D_ASSERT(ffi_expr);
+    auto &expr = reinterpret_cast<Expression *>(ffi_expr)->Cast<BoundCastExpression>();
+    return expr.try_cast;
+}
+
+extern "C" bool duckdb_vx_logical_type_contains_128bit(duckdb_logical_type ffi_type) {
+    D_ASSERT(ffi_type);
+    auto &type = *reinterpret_cast<LogicalType *>(ffi_type);
+    return TypeVisitor::Contains(type, [](const LogicalType &t) {
+        return t.id() == LogicalTypeId::HUGEINT || t.id() == LogicalTypeId::UHUGEINT;
+    });
+}

vortex-duckdb/cpp/include/cast_pushdown.hpp

@@ -0,0 +1,43 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+#pragma once
+#include "optimizer.hpp"
+
+#include "duckdb/common/unordered_set.hpp"
+#include "duckdb/main/client_context.hpp"
+#include "duckdb/planner/expression.hpp"
+#include "duckdb/planner/logical_operator.hpp"
+
+using namespace duckdb;
+
+/**
+ * Collect CAST(col) expressions. If "col" is used without CAST in "plan",
+ * record in "analyses.conflicts"
+ */
+struct CastCollect final : LogicalOperatorVisitor {
+    Analyses &analyses;
+    const Projections &projections;
+    // Casts that are direct outputs of a clean projection over a GET. Only these
+    // start a pushdown candidate; a nested cast may push down a different value.
+    // See test/sql/copy/csv/auto/test_large_integer_detection.test in duckdb
+    unordered_set<const Expression *> top_level_casts;
+
+    CastCollect(Analyses &analyses, const Projections &projections);
+    void VisitOperator(LogicalOperator &op) override;
+    ExpressionPtr VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) override;
+    ExpressionPtr VisitReplace(BoundCastExpression &expr, ExpressionPtr *ptr) override;
+};
+
+/*
+ * For "col" in columns collected by CastCollect, replace CAST(col) to "col"
+ * if "col" doesn't have conflicting usage. Update return types for bound
+ * columns and logical projections referencing this column.
+ */
+struct CastReplace final : LogicalOperatorVisitor {
+    Analyses &analyses;
+    const Projections &projections;
+
+    CastReplace(Analyses &analyses, const Projections &aliases);
+    ExpressionPtr VisitReplace(BoundColumnRefExpression &expr, ExpressionPtr *ptr) override;
+    ExpressionPtr VisitReplace(BoundCastExpression &expr, ExpressionPtr *ptr) override;
+};

vortex-duckdb/cpp/include/expr.h

@@ -264,6 +264,12 @@ typedef struct {
 
 void duckdb_vx_expr_get_bound_function(duckdb_vx_expr expr, duckdb_vx_expr_bound_function *out);
 
+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);
+
+bool duckdb_vx_logical_type_contains_128bit(duckdb_logical_type type);
+
 #ifdef __cplusplus /* End C ABI */
 }
 #endif