[GJ-98] Enable TPCH Q6 & Q1 for arrow backend

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

cpp/gazelle-cpp/compute/substrait_arrow.cc

@@ -17,6 +17,136 @@
 
 #include "substrait_arrow.h"
 
+#include <arrow/compute/exec/options.h>
+#include <arrow/compute/registry.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<gazellejni::RecordBatchResultIterator>
+ArrowExecBackend::GetResultIterator() {
+ return GetResultIterator({});
+}
+
+std::shared_ptr<gazellejni::RecordBatchResultIterator>
+ArrowExecBackend::GetResultIterator(
+ std::vector<std::shared_ptr<gazellejni::RecordBatchResultIterator>> 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<arrow::compute::Declaration>(std::move(decls[0]));
+
+ // Prepare and add source decls
+ if (!inputs.empty()) {
+ std::vector<arrow::compute::Declaration> 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<arrow::RecordBatch>& 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<arrow::util::optional<arrow::compute::ExecBatch>> 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<arrow::RecordBatchReader> sink_reader =
+ arrow::compute::MakeGeneratorReader(std::move(output_schema), std::move(sink_gen),
+ arrow::default_memory_pool());
+ return std::make_shared<gazellejni::RecordBatchResultIterator>(std::move(sink_reader),
+ shared_from_this());
+}
+
+void ArrowExecBackend::ReplaceSourceDecls(
+ std::vector<arrow::compute::Declaration> source_decls) {
+ std::vector<arrow::compute::Declaration*> visited;
+ std::vector<arrow::compute::Declaration*> 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<arrow::compute::Declaration>(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<arrow::compute::SourceIndexOptions>(
+ 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

cpp/gazelle-cpp/compute/substrait_arrow.h

@@ -17,30 +17,33 @@
 
 #pragma once
 
+#include <arrow/engine/substrait/serde.h>
+
 #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<gazellejni::compute::SubstraitParser>();
-}
+ArrowExecBackend() = default;
+
+~ArrowExecBackend() override;
 
- std::shared_ptr<gazellejni::RecordBatchResultIterator> GetResultIterator() override {
- return delegate_->GetResultIterator();
- }
+ std::shared_ptr<gazellejni::RecordBatchResultIterator> GetResultIterator() override;
 
  std::shared_ptr<gazellejni::RecordBatchResultIterator> GetResultIterator(
  std::vector<std::shared_ptr<gazellejni::RecordBatchResultIterator>> inputs)
- override {
- return delegate_->GetResultIterator(std::move(inputs));
- }
+ override;
 
  private:
- std::unique_ptr<gazellejni::compute::SubstraitParser> delegate_;
+ std::shared_ptr<arrow::compute::Declaration> decl_;
+ std::shared_ptr<arrow::compute::ExecPlan> exec_plan_;
+
+ void ReplaceSourceDecls(std::vector<arrow::compute::Declaration> source_decls);
 };
 
+void Initialize();
+
 } // namespace compute
 } // namespace gazellecpp

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<gazellecpp::compute::ArrowSubstraitParser>(); });
+ [] { return std::make_shared<gazellecpp::compute::ArrowExecBackend>(); });
 }
 
 #ifdef __cplusplus

cpp/src/jni/exec_backend.h

@@ -52,22 +52,27 @@ class RecordBatchResultIterator : public ResultIteratorBase<arrow::RecordBatch>
  template <typename T>
  explicit RecordBatchResultIterator(std::shared_ptr<T> iter,
  std::shared_ptr<ExecBackendBase> backend = nullptr)
- : iter_(std::make_shared<arrow::RecordBatchIterator>(Wrapper<T>(std::move(iter)))),
+ : iter_(std::make_unique<arrow::RecordBatchIterator>(Wrapper<T>(std::move(iter)))),
  next_(nullptr),
  backend_(std::move(backend)) {}
 
  bool HasNext() override {
+ CheckValid();
  GetNext();
  return next_ != nullptr;
  }
 
  std::shared_ptr<arrow::RecordBatch> 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<arrow::RecordBatchIterator> ToArrowRecordBatchIterator() {
- return iter_;
+ return std::move(iter_);
  }
 
  private:
@@ -82,11 +87,18 @@ class RecordBatchResultIterator : public ResultIteratorBase<arrow::RecordBatch>
  std::shared_ptr<T> ptr_;
  };
 
- std::shared_ptr<arrow::RecordBatchIterator> iter_;
+ std::unique_ptr<arrow::RecordBatchIterator> iter_;
  std::shared_ptr<arrow::RecordBatch> next_;
  std::shared_ptr<ExecBackendBase> 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<ExecBackendBase> {
  }
 
  /// Parse and get the input schema from the cached plan.
- arrow::Status GetInputSchemaMap(
- std::unordered_map<uint64_t, std::shared_ptr<arrow::Schema>>& 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<uint64_t, std::shared_ptr<arrow::Schema>>&
+ 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<ExecBackendBase> {
 
  protected:
  substrait::Plan plan_;
+ std::unordered_map<uint64_t, std::shared_ptr<arrow::Schema>> schema_map_;
 
  arrow::Result<std::shared_ptr<arrow::DataType>> subTypeToArrowType(
  const substrait::Type& stype) {
@@ -162,25 +179,23 @@ class ExecBackendBase : public std::enable_shared_from_this<ExecBackendBase> {
  case substrait::Type::KindCase::kString:
  return arrow::utf8();
  default:
- return arrow::Result<std::shared_ptr<arrow::DataType>>(
-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<uint64_t, std::shared_ptr<arrow::Schema>>& 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<ExecBackendBase> {
  }
 
  // 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<ExecBackendBase> {
  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();
  }
 };