Cast pushdown for duckdb

vortex-array/src/arrays/bool/compute/cast.rs

@@ -1,6 +1,9 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use num_traits::One;
+use num_traits::Zero;
+use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 
 use crate::ArrayRef;
@@ -9,8 +12,10 @@ use crate::IntoArray;
 use crate::array::ArrayView;
 use crate::arrays::Bool;
 use crate::arrays::BoolArray;
+use crate::arrays::PrimitiveArray;
 use crate::arrays::bool::BoolArrayExt;
 use crate::dtype::DType;
+use crate::match_each_native_ptype;
 use crate::scalar_fn::fns::cast::CastKernel;
 use crate::scalar_fn::fns::cast::CastReduce;
 
@@ -38,17 +43,34 @@ impl CastKernel for Bool {
         dtype: &DType,
         ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
-        if !dtype.is_boolean() {
-            return Ok(None);
+        if dtype.is_boolean() {
+            let new_validity =
+                array
+                    .validity()?
+                    .cast_nullability(dtype.nullability(), array.len(), ctx)?;
+            return Ok(Some(
+                BoolArray::new(array.to_bit_buffer(), new_validity).into_array(),
+            ));
         }
 
+        let DType::Primitive(new_ptype, new_nullability) = dtype else {
+            return Ok(None);
+        };
+
         let new_validity =
             array
                 .validity()?
-                .cast_nullability(dtype.nullability(), array.len(), ctx)?;
-        Ok(Some(
-            BoolArray::new(array.to_bit_buffer(), new_validity).into_array(),
-        ))
+                .cast_nullability(*new_nullability, array.len(), ctx)?;
+
+        let bits = array.to_bit_buffer();
+        let len = bits.len();
+
+        Ok(Some(match_each_native_ptype!(*new_ptype, |T| {
+            let (one, zero) = (<T as One>::one(), <T as Zero>::zero());
+            let mut buffer = BufferMut::<T>::with_capacity(len);
+            buffer.extend(bits.iter().map(|v| if v { one } else { zero }));
+            PrimitiveArray::new(buffer.freeze(), new_validity).into_array()
+        })))
     }
 }
 
@@ -102,4 +124,22 @@ mod tests {
     fn test_cast_bool_conformance(#[case] array: BoolArray) {
         test_cast_conformance(&array.into_array());
     }
+
+    #[rstest]
+    #[case(crate::dtype::PType::I8)]
+    #[case(crate::dtype::PType::I32)]
+    #[case(crate::dtype::PType::I64)]
+    #[case(crate::dtype::PType::U8)]
+    #[case(crate::dtype::PType::U64)]
+    #[case(crate::dtype::PType::F32)]
+    #[case(crate::dtype::PType::F64)]
+    fn cast_bool_to_primitive(#[case] target: crate::dtype::PType) {
+        let mut ctx = SESSION.create_execution_ctx();
+        let arr = BoolArray::from_iter(vec![true, false, true]).into_array();
+        let out = arr
+            .cast(DType::Primitive(target, Nullability::NonNullable))
+            .unwrap();
+        let out = out.execute::<Canonical>(&mut ctx).unwrap().into_array();
+        assert_eq!(out.len(), 3);
+    }
 }

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,169 @@
+// 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) {
+        // 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,17 @@ 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_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 ScalarFnCollect, 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);
+
+// Check if type or contained types i.e. List(T) contains HUGEINT/UHUGEINT
+// These are not present in DType so we can't convert.
+bool duckdb_vx_logical_type_contains_128bit(duckdb_logical_type type);
+
 #ifdef __cplusplus /* End C ABI */
 }
 #endif