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..06f3973dab2 100644 --- a/cpp/gazelle-cpp/compute/substrait_arrow.cc +++ b/cpp/gazelle-cpp/compute/substrait_arrow.cc @@ -17,6 +17,136 @@ #include "substrait_arrow.h" +#include +#include + +#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())); + } + decl_ = std::make_shared(std::move(decls[0])); + + // Prepare and add source decls + if (!inputs.empty()) { + std::vector 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)}); + } + 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())); + 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; + std::cout << "Execplan output schema:" << std::endl + << output_schema->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::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); + } + } + } + + 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; + } +} + +} // namespace compute } // namespace gazellecpp diff --git a/cpp/gazelle-cpp/compute/substrait_arrow.h b/cpp/gazelle-cpp/compute/substrait_arrow.h index 0c320f9f7b1..93c77c4b31d 100644 --- a/cpp/gazelle-cpp/compute/substrait_arrow.h +++ b/cpp/gazelle-cpp/compute/substrait_arrow.h @@ -17,30 +17,33 @@ #pragma once +#include + #include "compute/substrait_utils.h" namespace gazellecpp { namespace compute { -class ArrowSubstraitParser : public gazellejni::ExecBackendBase { +class ArrowExecBackend : public gazellejni::ExecBackendBase { public: - ArrowSubstraitParser() { - delegate_ = std::make_unique(); - } + ArrowExecBackend() = default; + + ~ArrowExecBackend() override; - std::shared_ptr GetResultIterator() override { - return delegate_->GetResultIterator(); - } + 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 decl_; + std::shared_ptr exec_plan_; + + 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 84a3d2e09ad..0913d98cdc3 100644 --- a/cpp/gazelle-cpp/jni/jni_wrapper.cc +++ b/cpp/gazelle-cpp/jni/jni_wrapper.cc @@ -43,8 +43,9 @@ 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(); }); + [] { return std::make_shared(); }); } #ifdef __cplusplus diff --git a/cpp/src/jni/exec_backend.h b/cpp/src/jni/exec_backend.h index 187f3280d18..1d6731fd90b 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) { @@ -162,25 +179,23 @@ 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())); } } 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()