From e113c16af95dc4392aebaa3ff2f094177f34fe68 Mon Sep 17 00:00:00 2001 From: Rong Ma Date: Wed, 16 Mar 2022 12:34:32 +0800 Subject: [PATCH 1/2] Q6 --- cpp/gazelle-cpp/CMakeLists.txt | 25 +++++- cpp/gazelle-cpp/compute/substrait_arrow.cc | 99 +++++++++++++++++++++- cpp/gazelle-cpp/compute/substrait_arrow.h | 25 +++--- cpp/gazelle-cpp/jni/jni_wrapper.cc | 2 +- cpp/src/jni/exec_backend.h | 71 ++++++++++------ cpp/src/jni/jni_wrapper.cc | 22 ++--- 6 files changed, 191 insertions(+), 53 deletions(-) diff --git a/cpp/gazelle-cpp/CMakeLists.txt b/cpp/gazelle-cpp/CMakeLists.txt index acd3b999b1e..d086b43239c 100644 --- a/cpp/gazelle-cpp/CMakeLists.txt +++ b/cpp/gazelle-cpp/CMakeLists.txt @@ -29,14 +29,37 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) find_package(JNI REQUIRED) +set(ARROW_ENGINE_LIB_NAME "arrow_engine") +set(ARROW_SHARED_LIBRARY_SUFFIX ".so.700") +find_library(ARROW_ENGINE_LIB NAMES ${CMAKE_SHARED_LIBRARY_PREFIX}${ARROW_ENGINE_LIB_NAME}${ARROW_SHARED_LIBRARY_SUFFIX}) +if (NOT ARROW_ENGINE_LIB) + message(FATAL_ERROR "Arrow Engine Library Not Found") +else () + message(STATUS "Arrow Engine Library Can Be Found in ${ARROW_ENGINE_LIB}") +endif () +file(COPY ${ARROW_ENGINE_LIB} DESTINATION ${root_directory}/releases/ FOLLOW_SYMLINK_CHAIN) + +# Set up Arrow Engine Shared Library Directory +set( + ARROW_ENGINE_SHARED_LIB + "${root_directory}/releases/${CMAKE_SHARED_LIBRARY_PREFIX}${ARROW_ENGINE_LIB_NAME}${ARROW_SHARED_LIBRARY_SUFFIX}" +) + +add_library(Arrow::arrow_engine SHARED IMPORTED) +set_target_properties(Arrow::arrow_engine + PROPERTIES IMPORTED_LOCATION "${ARROW_ENGINE_SHARED_LIB}" + INTERFACE_INCLUDE_DIRECTORIES + "${root_directory}/releases/include") + set(GAZELLE_CPP_JNI_SRCS jni/jni_wrapper.cc + compute/substrait_arrow.cc ) add_library(gazelle_cpp SHARED ${GAZELLE_CPP_JNI_SRCS}) target_include_directories(gazelle_cpp PUBLIC ${CMAKE_SYSTEM_INCLUDE_PATH} ${JNI_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR} ${root_directory}/src) -target_link_libraries(gazelle_cpp spark_columnar_jni) +target_link_libraries(gazelle_cpp spark_columnar_jni Arrow::arrow_engine) set_target_properties(gazelle_cpp PROPERTIES LIBRARY_OUTPUT_DIRECTORY ${root_directory}/releases diff --git a/cpp/gazelle-cpp/compute/substrait_arrow.cc b/cpp/gazelle-cpp/compute/substrait_arrow.cc index d654f330523..50029830626 100644 --- a/cpp/gazelle-cpp/compute/substrait_arrow.cc +++ b/cpp/gazelle-cpp/compute/substrait_arrow.cc @@ -17,6 +17,103 @@ #include "substrait_arrow.h" +#include "jni/exec_backend.h" + namespace gazellecpp { -namespace compute {} // namespace compute +namespace compute { + +ArrowExecBackend::~ArrowExecBackend() { + if (exec_plan_ != nullptr) { + exec_plan_->finished().Wait(); + } +#ifdef DEBUG + std::cout << "Plan finished" << std::endl; +#endif +} + +std::shared_ptr +ArrowExecBackend::GetResultIterator() { + return GetResultIterator({}); +} + +std::shared_ptr +ArrowExecBackend::GetResultIterator( + std::vector> inputs) { + GAZELLE_JNI_ASSIGN_OR_THROW(auto decls, arrow::engine::ConvertPlan(plan_)); + if (decls.size() != 1) { + throw gazellejni::JniPendingException("Expected 1 decl, but got " + + std::to_string(decls.size())); + } + auto& decl = decls[0]; + + // Prepare and add source decls + if (!inputs.empty()) { + std::deque source_decls; + for (auto i = 0; i < inputs.size(); ++i) { + auto it = schema_map_.find(i); + if (it == schema_map_.end()) { + throw gazellejni::JniPendingException( + "Schema not found for input batch iterator " + std::to_string(i)); + } + auto batch_it = MakeMapIterator( + [](const std::shared_ptr& batch) { + return arrow::util::make_optional(arrow::compute::ExecBatch(*batch)); + }, + std::move(*inputs[i]->ToArrowRecordBatchIterator())); + GAZELLE_JNI_ASSIGN_OR_THROW( + auto gen, arrow::MakeBackgroundGenerator(std::move(batch_it), + arrow::internal::GetCpuThreadPool())); + source_decls.emplace_back( + "source", arrow::compute::SourceNodeOptions{it->second, std::move(gen)}); + } + AddSourceDecls(decl, source_decls); + } + + // Make plan + GAZELLE_JNI_ASSIGN_OR_THROW(exec_plan_, arrow::compute::ExecPlan::Make()); + GAZELLE_JNI_ASSIGN_OR_THROW(auto node, decl.AddToPlan(exec_plan_.get())); + auto output_schema = node->output_schema(); + + // Add sink node. It's added after constructing plan from decls because sink node + // doesn't have output schema. + arrow::AsyncGenerator> sink_gen; + GAZELLE_JNI_THROW_NOT_OK(arrow::compute::MakeExecNode( + "sink", exec_plan_.get(), {node}, arrow::compute::SinkNodeOptions{&sink_gen})); + + GAZELLE_JNI_THROW_NOT_OK(exec_plan_->Validate()); + GAZELLE_JNI_THROW_NOT_OK(exec_plan_->StartProducing()); + +#ifdef DEBUG + std::cout << std::string(50, '#') << " produced arrow::ExecPlan:" << std::endl; + std::cout << exec_plan_->ToString() << std::endl; +#endif + + std::shared_ptr sink_reader = + arrow::compute::MakeGeneratorReader(std::move(output_schema), std::move(sink_gen), + arrow::default_memory_pool()); + return std::make_shared(std::move(sink_reader), + shared_from_this()); +} + +void ArrowExecBackend::AddSourceDecls( + arrow::compute::Declaration& decl, + std::deque& source_decls) { + if (decl.inputs.empty()) { + auto need_input = std::find(no_inputs.begin(), no_inputs.end(), decl.factory_name) == + no_inputs.end(); + if (need_input && !source_decls.empty()) { + decl.inputs.emplace_back(std::move(source_decls.front())); + source_decls.pop_front(); + } + return; + } + for (auto& input : decl.inputs) { + AddSourceDecls(arrow::util::get(input), source_decls); + if (source_decls.empty()) { + return; + } + } +} + +} // namespace compute } // namespace gazellecpp diff --git a/cpp/gazelle-cpp/compute/substrait_arrow.h b/cpp/gazelle-cpp/compute/substrait_arrow.h index 0c320f9f7b1..c51fcaf8da7 100644 --- a/cpp/gazelle-cpp/compute/substrait_arrow.h +++ b/cpp/gazelle-cpp/compute/substrait_arrow.h @@ -17,29 +17,32 @@ #pragma once +#include + #include "compute/substrait_utils.h" namespace gazellecpp { namespace compute { -class ArrowSubstraitParser : public gazellejni::ExecBackendBase { +static const std::vector no_inputs = {"scan", "source", "table_source"}; + +class ArrowExecBackend : public gazellejni::ExecBackendBase { public: - ArrowSubstraitParser() { - delegate_ = std::make_unique(); - } + ArrowExecBackend() = default; - std::shared_ptr GetResultIterator() override { - return delegate_->GetResultIterator(); - } + ~ArrowExecBackend() override; + + std::shared_ptr GetResultIterator() override; std::shared_ptr GetResultIterator( std::vector> inputs) - override { - return delegate_->GetResultIterator(std::move(inputs)); - } + override; private: - std::unique_ptr delegate_; + std::shared_ptr exec_plan_; + + void AddSourceDecls(arrow::compute::Declaration& decl, + std::deque& source_decls); }; } // namespace compute diff --git a/cpp/gazelle-cpp/jni/jni_wrapper.cc b/cpp/gazelle-cpp/jni/jni_wrapper.cc index 84a3d2e09ad..362e08be3aa 100644 --- a/cpp/gazelle-cpp/jni/jni_wrapper.cc +++ b/cpp/gazelle-cpp/jni/jni_wrapper.cc @@ -44,7 +44,7 @@ JNIEXPORT void JNICALL Java_com_intel_oap_vectorized_ExpressionEvaluatorJniWrapper_nativeInitNative( JNIEnv* env, jobject obj) { gazellejni::SetBackendFactory( - [] { return std::make_shared(); }); + [] { return std::make_shared(); }); } #ifdef __cplusplus diff --git a/cpp/src/jni/exec_backend.h b/cpp/src/jni/exec_backend.h index 187f3280d18..132e548fcc6 100644 --- a/cpp/src/jni/exec_backend.h +++ b/cpp/src/jni/exec_backend.h @@ -52,22 +52,27 @@ class RecordBatchResultIterator : public ResultIteratorBase template explicit RecordBatchResultIterator(std::shared_ptr iter, std::shared_ptr backend = nullptr) - : iter_(std::make_shared(Wrapper(std::move(iter)))), + : iter_(std::make_unique(Wrapper(std::move(iter)))), next_(nullptr), backend_(std::move(backend)) {} bool HasNext() override { + CheckValid(); GetNext(); return next_ != nullptr; } std::shared_ptr Next() override { + CheckValid(); GetNext(); return std::move(next_); } + /// arrow::RecordBatchIterator doesn't support shared ownership. Once this method is + /// called, the caller should take it's ownership, and RecordBatchResultIterator + /// will no longer have access to the underlying iterator. std::shared_ptr ToArrowRecordBatchIterator() { - return iter_; + return std::move(iter_); } private: @@ -82,11 +87,18 @@ class RecordBatchResultIterator : public ResultIteratorBase std::shared_ptr ptr_; }; - std::shared_ptr iter_; + std::unique_ptr iter_; std::shared_ptr next_; std::shared_ptr backend_; - void GetNext() { + inline void CheckValid() { + if (iter_ == nullptr) { + throw JniPendingException( + "RecordBatchResultIterator: the underlying iterator has expired."); + } + } + + inline void GetNext() { if (next_ == nullptr) { GAZELLE_JNI_ASSIGN_OR_THROW(next_, iter_->Next()); } @@ -117,22 +129,26 @@ class ExecBackendBase : public std::enable_shared_from_this { } /// Parse and get the input schema from the cached plan. - arrow::Status GetInputSchemaMap( - std::unordered_map>& schema_map) { - for (auto& srel : plan_.relations()) { - if (srel.has_root()) { - auto& sroot = srel.root(); - if (sroot.has_input()) { - GetIterInputSchemaFromRel(sroot.input(), schema_map); - } else { - throw std::runtime_error("Expect Rel as input."); + const std::unordered_map>& + GetInputSchemaMap() { + if (schema_map_.empty()) { + for (auto& srel : plan_.relations()) { + if (srel.has_root()) { + auto& sroot = srel.root(); + if (sroot.has_input()) { + // TODO: remove arrow::Status + GAZELLE_JNI_THROW_NOT_OK(GetIterInputSchemaFromRel(sroot.input())); + } else { + throw JniPendingException("Expect Rel as input."); + } + } + if (srel.has_rel()) { + // TODO: remove arrow::Status + GAZELLE_JNI_THROW_NOT_OK(GetIterInputSchemaFromRel(srel.rel())); } - } - if (srel.has_rel()) { - GetIterInputSchemaFromRel(srel.rel(), schema_map); } } - return arrow::Status::OK(); + return schema_map_; } /// This function is used to create certain converter from the format used by the @@ -146,6 +162,7 @@ class ExecBackendBase : public std::enable_shared_from_this { protected: substrait::Plan plan_; + std::unordered_map> schema_map_; arrow::Result> subTypeToArrowType( const substrait::Type& stype) { @@ -168,19 +185,17 @@ class ExecBackendBase : public std::enable_shared_from_this { } private: - // This method is used to get the input schema in ReadRel. - arrow::Status GetIterInputSchemaFromRel( - const substrait::Rel& srel, - std::unordered_map>& schema_map) { + // This method is used to get the input schema in InputRel. + arrow::Status GetIterInputSchemaFromRel(const substrait::Rel& srel) { // TODO: need to support more Substrait Rels here. if (srel.has_aggregate() && srel.aggregate().has_input()) { - return GetIterInputSchemaFromRel(srel.aggregate().input(), schema_map); + return GetIterInputSchemaFromRel(srel.aggregate().input()); } if (srel.has_project() && srel.project().has_input()) { - return GetIterInputSchemaFromRel(srel.project().input(), schema_map); + return GetIterInputSchemaFromRel(srel.project().input()); } if (srel.has_filter() && srel.filter().has_input()) { - return GetIterInputSchemaFromRel(srel.filter().input(), schema_map); + return GetIterInputSchemaFromRel(srel.filter().input()); } if (!srel.has_read()) { return arrow::Status::Invalid("Read Rel expected."); @@ -216,7 +231,6 @@ class ExecBackendBase : public std::enable_shared_from_this { } // Get the iterator index. - int32_t iterIdx; if (sread.has_local_files()) { const auto& fileList = sread.local_files().items(); if (fileList.size() == 0) { @@ -229,11 +243,12 @@ class ExecBackendBase : public std::enable_shared_from_this { return arrow::Status::Invalid("Iterator index is not found."); } std::string idxStr = filePath.substr(pos + prefix.size(), filePath.size()); - iterIdx = std::stoi(idxStr); + auto iterIdx = std::stoi(idxStr); + + // Set up the schema map. + schema_map_[iterIdx] = arrow::schema(input_fields); } - // Set up the schema map. - schema_map[iterIdx] = arrow::schema(input_fields); return arrow::Status::OK(); } }; diff --git a/cpp/src/jni/jni_wrapper.cc b/cpp/src/jni/jni_wrapper.cc index 2ce071a55c2..2d26ab419ab 100644 --- a/cpp/src/jni/jni_wrapper.cc +++ b/cpp/src/jni/jni_wrapper.cc @@ -369,23 +369,21 @@ Java_com_intel_oap_vectorized_ExpressionEvaluatorJniWrapper_nativeCreateKernelWi jsize iters_len = env->GetArrayLength(iter_arr); std::vector> input_iters; if (iters_len > 0) { - // Construct a map between iter index and input schema. - std::unordered_map> schema_map; // Get input schema from Substrait plan. - backend->GetInputSchemaMap(schema_map); + const auto& schema_map = backend->GetInputSchemaMap(); for (int idx = 0; idx < iters_len; idx++) { jobject iter = env->GetObjectArrayElement(iter_arr, idx); // IMPORTANT: DO NOT USE LOCAL REF IN DIFFERENT THREAD // TODO Release this in JNI Unload or dependent object's destructor jobject ref_iter = env->NewGlobalRef(iter); - std::shared_ptr input_schema; - if (schema_map.find(idx) == schema_map.end()) { - ThrowPendingException("Not found the schema."); - } else { - input_schema = schema_map[idx]; + auto it = schema_map.find(idx); + if (it == schema_map.end()) { + ThrowPendingException("Schema not found for input batch iterator " + + std::to_string(idx)); } - auto rb_iter = std::make_shared(vm, ref_iter, - std::move(input_schema)); + + auto rb_iter = std::make_shared( + vm, ref_iter, it->second); input_iters.push_back( std::make_shared(std::move(rb_iter))); } @@ -431,8 +429,10 @@ JNIEXPORT void JNICALL Java_com_intel_oap_vectorized_BatchIterator_nativeClose( #ifdef DEBUG auto it = batch_iterator_holder_.Lookup(id); if (it.use_count() > 2) { - std::cout << "Id " << id << " use count is " << it.use_count() << std::endl; + std::cout << "RecordBatchResultIterator Id " << id << " use count is " + << it.use_count() << std::endl; } + std::cout << "BatchIterator nativeClose." << std::endl; #endif batch_iterator_holder_.Erase(id); JNI_METHOD_END() From 30bb1e513a28e39ba0c2e989dd06164dd52e4f16 Mon Sep 17 00:00:00 2001 From: Rong Ma Date: Wed, 6 Apr 2022 12:59:11 +0800 Subject: [PATCH 2/2] TPCH Q1 --- cpp/gazelle-cpp/compute/substrait_arrow.cc | 71 ++++++++++++++++------ cpp/gazelle-cpp/compute/substrait_arrow.h | 8 +-- cpp/gazelle-cpp/jni/jni_wrapper.cc | 1 + cpp/src/jni/exec_backend.h | 4 +- 4 files changed, 59 insertions(+), 25 deletions(-) diff --git a/cpp/gazelle-cpp/compute/substrait_arrow.cc b/cpp/gazelle-cpp/compute/substrait_arrow.cc index 50029830626..06f3973dab2 100644 --- a/cpp/gazelle-cpp/compute/substrait_arrow.cc +++ b/cpp/gazelle-cpp/compute/substrait_arrow.cc @@ -17,6 +17,9 @@ #include "substrait_arrow.h" +#include +#include + #include "jni/exec_backend.h" namespace gazellecpp { @@ -44,11 +47,11 @@ ArrowExecBackend::GetResultIterator( throw gazellejni::JniPendingException("Expected 1 decl, but got " + std::to_string(decls.size())); } - auto& decl = decls[0]; + decl_ = std::make_shared(std::move(decls[0])); // Prepare and add source decls if (!inputs.empty()) { - std::deque source_decls; + std::vector source_decls; for (auto i = 0; i < inputs.size(); ++i) { auto it = schema_map_.find(i); if (it == schema_map_.end()) { @@ -66,12 +69,12 @@ ArrowExecBackend::GetResultIterator( source_decls.emplace_back( "source", arrow::compute::SourceNodeOptions{it->second, std::move(gen)}); } - AddSourceDecls(decl, source_decls); + ReplaceSourceDecls(std::move(source_decls)); } // Make plan GAZELLE_JNI_ASSIGN_OR_THROW(exec_plan_, arrow::compute::ExecPlan::Make()); - GAZELLE_JNI_ASSIGN_OR_THROW(auto node, decl.AddToPlan(exec_plan_.get())); + GAZELLE_JNI_ASSIGN_OR_THROW(auto node, decl_->AddToPlan(exec_plan_.get())); auto output_schema = node->output_schema(); // Add sink node. It's added after constructing plan from decls because sink node @@ -86,6 +89,8 @@ ArrowExecBackend::GetResultIterator( #ifdef DEBUG std::cout << std::string(50, '#') << " produced arrow::ExecPlan:" << std::endl; std::cout << exec_plan_->ToString() << std::endl; + std::cout << "Execplan output schema:" << std::endl + << output_schema->ToString() << std::endl; #endif std::shared_ptr sink_reader = @@ -95,23 +100,51 @@ ArrowExecBackend::GetResultIterator( shared_from_this()); } -void ArrowExecBackend::AddSourceDecls( - arrow::compute::Declaration& decl, - std::deque& source_decls) { - if (decl.inputs.empty()) { - auto need_input = std::find(no_inputs.begin(), no_inputs.end(), decl.factory_name) == - no_inputs.end(); - if (need_input && !source_decls.empty()) { - decl.inputs.emplace_back(std::move(source_decls.front())); - source_decls.pop_front(); +void ArrowExecBackend::ReplaceSourceDecls( + std::vector source_decls) { + std::vector visited; + std::vector source_indexes; + + visited.push_back(decl_.get()); + + while (!visited.empty()) { + auto top = visited.back(); + visited.pop_back(); + for (auto& input : top->inputs) { + auto& input_decl = arrow::util::get(input); + if (input_decl.factory_name == "source_index") { + source_indexes.push_back(&input_decl); + } else { + visited.push_back(&input_decl); + } } - return; } - for (auto& input : decl.inputs) { - AddSourceDecls(arrow::util::get(input), source_decls); - if (source_decls.empty()) { - return; - } + + if (source_indexes.size() != source_decls.size()) { + throw gazellejni::JniPendingException( + "Wrong number of source declarations. " + std::to_string(source_indexes.size()) + + " source(s) needed by source declarations, but got " + + std::to_string(source_decls.size()) + " from input batches."); + } + + for (auto& source_index : source_indexes) { + auto index = + arrow::internal::checked_pointer_cast( + source_index->options) + ->index; + *source_index = std::move(source_decls[index]); + } +} + +void Initialize() { + static auto function_registry = arrow::compute::GetFunctionRegistry(); + static auto extension_registry = arrow::engine::default_extension_id_registry(); + if (function_registry && extension_registry) { + // TODO: Register customized functions to function_registry, and register the + // mapping from substrait function names to customized function names to + // extension_registry. + function_registry = nullptr; + extension_registry = nullptr; } } diff --git a/cpp/gazelle-cpp/compute/substrait_arrow.h b/cpp/gazelle-cpp/compute/substrait_arrow.h index c51fcaf8da7..93c77c4b31d 100644 --- a/cpp/gazelle-cpp/compute/substrait_arrow.h +++ b/cpp/gazelle-cpp/compute/substrait_arrow.h @@ -24,8 +24,6 @@ namespace gazellecpp { namespace compute { -static const std::vector no_inputs = {"scan", "source", "table_source"}; - class ArrowExecBackend : public gazellejni::ExecBackendBase { public: ArrowExecBackend() = default; @@ -39,11 +37,13 @@ class ArrowExecBackend : public gazellejni::ExecBackendBase { override; private: + std::shared_ptr decl_; std::shared_ptr exec_plan_; - void AddSourceDecls(arrow::compute::Declaration& decl, - std::deque& source_decls); + void ReplaceSourceDecls(std::vector source_decls); }; +void Initialize(); + } // namespace compute } // namespace gazellecpp diff --git a/cpp/gazelle-cpp/jni/jni_wrapper.cc b/cpp/gazelle-cpp/jni/jni_wrapper.cc index 362e08be3aa..0913d98cdc3 100644 --- a/cpp/gazelle-cpp/jni/jni_wrapper.cc +++ b/cpp/gazelle-cpp/jni/jni_wrapper.cc @@ -43,6 +43,7 @@ void JNI_OnUnload(JavaVM* vm, void* reserved) { JNIEXPORT void JNICALL Java_com_intel_oap_vectorized_ExpressionEvaluatorJniWrapper_nativeInitNative( JNIEnv* env, jobject obj) { + gazellecpp::compute::Initialize(); gazellejni::SetBackendFactory( [] { return std::make_shared(); }); } diff --git a/cpp/src/jni/exec_backend.h b/cpp/src/jni/exec_backend.h index 132e548fcc6..1d6731fd90b 100644 --- a/cpp/src/jni/exec_backend.h +++ b/cpp/src/jni/exec_backend.h @@ -179,8 +179,8 @@ class ExecBackendBase : public std::enable_shared_from_this { case substrait::Type::KindCase::kString: return arrow::utf8(); default: - return arrow::Result>( - arrow::Status::Invalid("Type not supported: " + stype.kind_case())); + return arrow::Status::Invalid("Type not supported: " + + std::to_string(stype.kind_case())); } }