feat(c++): support passing filter and projection expression in scan (#4224)

I found that direct binding doesn't seem too burdensome compared to
using protobuf serde. Performance-wise, using binding avoids the serde
cost but requires putting the Rust data structures on the heap using
`Box::new`, not sure whether this will be a problem if we have complex
nested expr. On the other hand, it will make things easier is if
vortex-expr can consume substrait.

Scalar, DType, and Expr support are not complete (e.g.,
decimal/date/time is not included), but you can create DType using
`from_arrow` and cast a Scalar from primitive types to another.

---------

Signed-off-by: Xinyu Zeng <[email protected]>

Author: XinyuZeng

.github/workflows/ci.yml

@@ -461,7 +461,7 @@ jobs:
         run: |
           mkdir -p build
           cd build
-          cmake -DVORTEX_ENABLE_TESTING=ON ..
+          cmake -DVORTEX_ENABLE_TESTING=ON -DVORTEX_ENABLE_ASAN=ON ..
           cmake --build . --parallel $(nproc)
           ctest -V
       - name: Build and run the example in release mode

Cargo.lock

@@ -18,6 +18,7 @@ endif()
 set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -Wall -Wextra -Wpedantic -Wno-dollar-in-identifier-extension")
 
 option(VORTEX_ENABLE_TESTING "Enable building test binary for vortex-cxx" OFF)
+option(VORTEX_ENABLE_ASAN "Enable address sanitizer" OFF)
 
 include(FetchContent)
 FetchContent_Declare(
@@ -27,7 +28,7 @@ FetchContent_Declare(
 )
 FetchContent_MakeAvailable(Corrosion)
 
-set(RUST_SOURCE_FILE lib.rs read.rs write.rs)
+set(RUST_SOURCE_FILE lib.rs)
 
 # Import Rust crate using Corrosion
 corrosion_import_crate(
@@ -43,11 +44,8 @@ FetchContent_Declare(nanoarrow
     URL "https://www.apache.org/dyn/closer.lua?action=download&filename=arrow/apache-arrow-nanoarrow-0.7.0/apache-arrow-nanoarrow-0.7.0.tar.gz")
 FetchContent_MakeAvailable(nanoarrow)
 
-# C++ source files
-set(CPP_SOURCE_FILE
-    cpp/src/write_options.cpp
-    cpp/src/file.cpp
-    cpp/src/scan.cpp
+file(GLOB_RECURSE CPP_SOURCE_FILE CONFIGURE_DEPENDS
+    "${CMAKE_CURRENT_SOURCE_DIR}/cpp/src/*.cpp"
 )
 
 # Public headers
@@ -62,17 +60,21 @@ set(CPP_PRIVATE_INCLUDE_DIRS
 
 # Create the main library combining C++ and Rust code
 add_library(vortex STATIC ${CPP_SOURCE_FILE})
-target_include_directories(vortex PUBLIC ${CPP_INCLUDE_DIRS})
+target_include_directories(vortex PUBLIC ${CPP_INCLUDE_DIRS} 
+    ${CMAKE_CURRENT_BINARY_DIR}/corrosion_generated/cxxbridge/vortex_cxx_bridge/include)
 target_include_directories(vortex PRIVATE 
     ${CPP_PRIVATE_INCLUDE_DIRS}
-    ${CMAKE_CURRENT_BINARY_DIR}/corrosion_generated/cxxbridge/vortex_cxx_bridge/include
-    ${CMAKE_CURRENT_BINARY_DIR}/corrosion_generated/cxxbridge/vortex_cxx_bridge/src
 )
 target_link_libraries(vortex 
     PUBLIC nanoarrow_static
     PRIVATE vortex_cxx_bridge
 )
 
+if (VORTEX_ENABLE_ASAN) 
+    target_compile_options(vortex PRIVATE -fsanitize=leak,address,undefined -fno-omit-frame-pointer -fno-common -O1)
+    target_link_options(vortex PRIVATE -fsanitize=leak,address,undefined)
+endif()
+
 # Tests
 if (VORTEX_ENABLE_TESTING)
     FetchContent_Declare(
@@ -96,6 +98,10 @@ if (VORTEX_ENABLE_TESTING)
         set(APPLE_LINK_FLAGS "-framework CoreFoundation -framework Security")
     endif()
     target_link_libraries(vortex_cxx_test PRIVATE vortex_cxx_bridge ${APPLE_LINK_FLAGS})
+    if (VORTEX_ENABLE_ASAN) 
+        target_compile_options(vortex_cxx_test PRIVATE -fsanitize=leak,address,undefined -fno-omit-frame-pointer -fno-common -O1)
+        target_link_options(vortex_cxx_test PRIVATE -fsanitize=leak,address,undefined)
+    endif()
     include(GoogleTest)
     gtest_discover_tests(vortex_cxx_test)
 endif()

vortex-cxx/CMakeLists.txt

@@ -21,10 +21,12 @@ workspace = true
 crate-type = ["staticlib"]
 
 [dependencies]
+anyhow = { workspace = true }
 arrow-array = { workspace = true, features = ["ffi"] }
 arrow-schema = { workspace = true }
 cxx = "1.0"
 futures = { workspace = true, features = ["thread-pool"] }
+paste = { workspace = true }
 prost = { workspace = true }
 take_mut = { workspace = true }
 tokio = { workspace = true, features = ["rt", "rt-multi-thread", "macros"] }

vortex-cxx/Cargo.toml

@@ -6,4 +6,6 @@
 #include "vortex/file.hpp"
 #include "vortex/scan.hpp"
 #include "vortex/write_options.hpp"
-#include "vortex/exception.hpp"
+#include "vortex/exception.hpp"
+#include "vortex/expr.hpp"
+#include "vortex/scalar.hpp"

vortex-cxx/cpp/include/vortex.hpp

@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#pragma once
+
+#include <string>
+#include <nanoarrow/common/inline_types.h>
+#include "vortex_cxx_bridge/lib.h"
+
+namespace vortex {
+
+enum class PType : uint8_t {
+    U8 = 0,
+    U16,
+    U32,
+    U64,
+    I8,
+    I16,
+    I32,
+    I64,
+    F16,
+    F32,
+    F64,
+};
+
+namespace dtype {
+    class DType {
+    public:
+        DType() = delete;
+        explicit DType(rust::Box<ffi::DType> impl) : impl_(std::move(impl)) {
+        }
+        DType(DType &&other) noexcept = default;
+        DType &operator=(DType &&other) = default;
+        ~DType() = default;
+
+        DType(const DType &) = delete;
+        DType &operator=(const DType &) = delete;
+
+        std::string ToString() const;
+
+        const rust::Box<ffi::DType> &GetImpl() {
+            return impl_;
+        }
+
+    private:
+        rust::Box<ffi::DType> impl_;
+    };
+
+    // Factory functions
+    DType null();
+    DType bool_(bool nullable = false);
+    DType primitive(PType ptype, bool nullable = false);
+    DType int8(bool nullable = false);
+    DType int16(bool nullable = false);
+    DType int32(bool nullable = false);
+    DType int64(bool nullable = false);
+    DType uint8(bool nullable = false);
+    DType uint16(bool nullable = false);
+    DType uint32(bool nullable = false);
+    DType uint64(bool nullable = false);
+    DType float16(bool nullable = false);
+    DType float32(bool nullable = false);
+    DType float64(bool nullable = false);
+    DType decimal(uint8_t precision = 10, int8_t scale = 0, bool nullable = false);
+    DType utf8(bool nullable = false);
+    DType binary(bool nullable = false);
+    /// TODO: Other DTypes are only supported by creating from Arrow for now.
+    DType from_arrow(struct ArrowSchema &schema, bool non_nullable = false);
+} // namespace dtype
+
+} // namespace vortex

vortex-cxx/cpp/include/vortex/dtype.hpp

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#pragma once
+
+#include "vortex/scalar.hpp"
+#include "vortex_cxx_bridge/lib.h"
+#include <vector>
+
+namespace vortex::expr {
+using scalar::Scalar;
+class Expr {
+public:
+    Expr() = delete;
+    explicit Expr(rust::Box<ffi::Expr> impl) : impl_(std::move(impl)) {
+    }
+    Expr(Expr &&other) noexcept = default;
+    Expr &operator=(Expr &&other) noexcept = default;
+    ~Expr() = default;
+
+    Expr(const Expr &) = delete;
+    Expr &operator=(const Expr &) = delete;
+
+    rust::Box<ffi::Expr> IntoImpl() && {
+        return std::move(impl_);
+    }
+
+private:
+    rust::Box<ffi::Expr> impl_;
+};
+
+Expr literal(Scalar scalar);
+Expr root();
+Expr column(std::string_view name);
+Expr get_item(std::string_view field, Expr expr);
+Expr not_(Expr expr);
+Expr is_null(Expr expr);
+Expr eq(Expr lhs, Expr rhs);
+Expr not_eq_(Expr lhs, Expr rhs);
+Expr gt(Expr lhs, Expr rhs);
+Expr gt_eq(Expr lhs, Expr rhs);
+Expr lt(Expr lhs, Expr rhs);
+Expr lt_eq(Expr lhs, Expr rhs);
+Expr and_(Expr lhs, Expr rhs);
+Expr or_(Expr lhs, Expr rhs);
+Expr checked_add(Expr lhs, Expr rhs);
+Expr select(const std::vector<std::string_view> &fields, Expr child);
+} // namespace vortex::expr

vortex-cxx/cpp/include/vortex/expr.hpp

@@ -4,8 +4,7 @@
 #pragma once
 
 #include <nanoarrow/common/inline_types.h>
-
-#include <memory>
+#include "vortex_cxx_bridge/lib.h"
 
 namespace vortex {
 class ScanBuilder;
@@ -14,9 +13,9 @@ class VortexFile {
 public:
     static VortexFile Open(const std::string &path);
 
-    VortexFile(VortexFile &&other) noexcept;
-    VortexFile &operator=(VortexFile &&other) noexcept;
-    ~VortexFile();
+    VortexFile(VortexFile &&other) noexcept = default;
+    VortexFile &operator=(VortexFile &&other) noexcept = default;
+    ~VortexFile() = default;
 
     VortexFile(const VortexFile &) = delete;
     VortexFile &operator=(const VortexFile &) = delete;
@@ -29,10 +28,10 @@ class VortexFile {
     ScanBuilder CreateScanBuilder() const;
 
 private:
-    struct Impl;
-    explicit VortexFile(std::unique_ptr<Impl> impl);
+    explicit VortexFile(rust::Box<ffi::VortexFile> impl) : impl_(std::move(impl)) {
+    }
 
-    std::unique_ptr<Impl> impl_;
+    rust::Box<ffi::VortexFile> impl_;
 };
 
 } // namespace vortex

vortex-cxx/cpp/include/vortex/file.hpp

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#pragma once
+
+#include <string_view>
+#include "dtype.hpp"
+#include "vortex_cxx_bridge/lib.h"
+
+namespace vortex::scalar {
+using dtype::DType;
+class Scalar {
+public:
+    Scalar() = delete;
+    explicit Scalar(rust::Box<ffi::Scalar> impl) : impl_(std::move(impl)) {
+    }
+    Scalar(Scalar &&other) noexcept = default;
+    Scalar &operator=(Scalar &&other) noexcept = default;
+    ~Scalar() = default;
+
+    Scalar(const Scalar &) = delete;
+    Scalar &operator=(const Scalar &) = delete;
+
+    rust::Box<ffi::Scalar> IntoImpl() && {
+        return std::move(impl_);
+    }
+
+private:
+    rust::Box<ffi::Scalar> impl_;
+};
+
+// Factory functions for creating scalar values
+Scalar bool_(bool value);
+Scalar int8(int8_t value);
+Scalar int16(int16_t value);
+Scalar int32(int32_t value);
+Scalar int64(int64_t value);
+Scalar uint8(uint8_t value);
+Scalar uint16(uint16_t value);
+Scalar uint32(uint32_t value);
+Scalar uint64(uint64_t value);
+Scalar float32(float value);
+Scalar float64(double value);
+Scalar string(std::string_view value);
+Scalar binary(const uint8_t *data, size_t length);
+/// TODO: Other Scalars are only supported by casting for now.
+Scalar cast(Scalar scalar, DType dtype);
+} // namespace vortex::scalar