[C++][Gandiva] Selectively register external C functions based on expression usage

Author: Reranko05

cpp/src/gandiva/engine.cc

@@ -348,10 +348,19 @@ Engine::~Engine() {}
 
 Status Engine::Init(std::unordered_set<std::string> function_names) {
   used_functions_ = std::move(function_names);
+  selective_mapping_enabled_ = true;
   std::call_once(register_exported_funcs_flag, gandiva::RegisterExportedFuncs);
 
   // Add mappings for global functions that can be accessed from LLVM/IR module.
   ARROW_RETURN_NOT_OK(AddGlobalMappings());
+  selective_mapping_enabled_ = false;
+  used_functions_.clear();
+  return Status::OK();
+}
+
+Status Engine::Init() {
+  std::call_once(register_exported_funcs_flag, gandiva::RegisterExportedFuncs);
+  ARROW_RETURN_NOT_OK(AddGlobalMappings());
   return Status::OK();
 }
 
@@ -395,6 +404,7 @@ Result<std::unique_ptr<Engine>> Engine::Make(
   std::unique_ptr<Engine> engine{
       new Engine(conf, std::move(jit), std::move(shared_target_machine), cached)};
 
+  ARROW_RETURN_NOT_OK(engine->Init());
   return engine;
 }
 
@@ -599,13 +609,6 @@ Result<void*> Engine::CompiledFunction(const std::string& function) {
 
 void Engine::AddGlobalMappingForFunc(const std::string& name, llvm::Type* ret_type,
                                      const std::vector<llvm::Type*>& args, void* func) {
-  bool is_internal_func =
-      internal_functions_.find(name) != internal_functions_.end();
-
-  if (!(is_internal_func ||
-        used_functions_.find(name) != used_functions_.end())) {
-    return;
-  }
   const auto prototype = llvm::FunctionType::get(ret_type, args, /*is_var_arg*/ false);
   llvm::Function::Create(prototype, llvm::GlobalValue::ExternalLinkage, name, module());
   AddAbsoluteSymbol(*lljit_, name, func);

cpp/src/gandiva/engine.h

@@ -94,8 +94,11 @@ class GANDIVA_EXPORT Engine {
   llvm::Constant* CreateGlobalStringPtr(const std::string& string);
 
   Status Init(std::unordered_set<std::string> function_names);
+  Status Init();
 
  private:
+  friend class ExternalCFunctions;
+  
   Engine(const std::shared_ptr<Configuration>& conf,
          std::unique_ptr<llvm::orc::LLJIT> lljit,
          std::shared_ptr<llvm::TargetMachine> target_machine, bool cached);
@@ -127,27 +130,28 @@ class GANDIVA_EXPORT Engine {
   std::unordered_set<std::string> used_functions_;
 
   static inline const std::unordered_set<std::string> internal_functions_ = {
-    "gdv_fn_context_arena_malloc",
-    "gdv_fn_context_set_error_msg",
-    "gdv_fn_populate_varlen_vector",
-    "gdv_fn_context_arena_reset",
-    "gdv_fn_in_expr_lookup_int32",
-    "gdv_fn_in_expr_lookup_int64",
-    "gdv_fn_in_expr_lookup_float",
-    "gdv_fn_in_expr_lookup_double",
-    "gdv_fn_in_expr_lookup_decimal",
-    "gdv_fn_in_expr_lookup_utf8",
-
-    "bitMapGetBit",
-    "bitMapSetBit",
-    "bitMapValidityGetBit",
-    "bitMapClearBitIfFalse",
+      "gdv_fn_context_arena_malloc",
+      "gdv_fn_context_set_error_msg",
+      "gdv_fn_populate_varlen_vector",
+      "gdv_fn_context_arena_reset",
+      "gdv_fn_in_expr_lookup_int32",
+      "gdv_fn_in_expr_lookup_int64",
+      "gdv_fn_in_expr_lookup_float",
+      "gdv_fn_in_expr_lookup_double",
+      "gdv_fn_in_expr_lookup_decimal",
+      "gdv_fn_in_expr_lookup_utf8",
+
+      "bitMapGetBit",
+      "bitMapSetBit",
+      "bitMapValidityGetBit",
+      "bitMapClearBitIfFalse",
   };
 
   bool optimize_ = true;
   bool module_finalized_ = false;
   bool cached_;
   bool functions_loaded_ = false;
+  bool selective_mapping_enabled_ = false;
   std::shared_ptr<FunctionRegistry> function_registry_;
   std::string module_ir_;
   // The lifetime of the TargetMachine is shared with LLJIT. This prevents unnecessary

cpp/src/gandiva/external_c_functions.cc

@@ -67,7 +67,21 @@ namespace gandiva {
 Status ExternalCFunctions::AddMappings(Engine* engine) const {
   const auto& c_funcs = function_registry_->GetCFunctions();
   const auto types = engine->types();
+
+  // Build allowed set ONCE before the loop
+  std::unordered_set<std::string> allowed;
+  if (engine->selective_mapping_enabled_) {
+    allowed = engine->internal_functions_;
+    allowed.insert(engine->used_functions_.begin(), engine->used_functions_.end());
+  }
+
   for (auto& [func, func_ptr] : c_funcs) {
+    const std::string& name = func.pc_name();
+
+    if (engine->selective_mapping_enabled_ && !allowed.contains(name)) {
+      continue;
+    }
+
     for (const auto& sig : func.signatures()) {
       ARROW_ASSIGN_OR_RAISE(auto llvm_signature, MapToLLVMSignature(sig, func, types));
       auto& [args, ret_llvm_type] = llvm_signature;

cpp/src/gandiva/llvm_generator.cc

@@ -67,23 +67,20 @@ Status LLVMGenerator::SetLLVMObjectCache(GandivaObjectCache& object_cache) {
   return engine_->SetLLVMObjectCache(object_cache);
 }
 
-arrow::Result<ValueValidityPairPtr> LLVMGenerator::Decompose(
-    const ExpressionPtr& expr) {
+arrow::Result<ValueValidityPairPtr> LLVMGenerator::Decompose(const ExpressionPtr& expr) {
   ExprDecomposer decomposer(*function_registry_, annotator_);
 
   ValueValidityPairPtr value_validity;
-  ARROW_RETURN_NOT_OK(
-      decomposer.Decompose(*expr->root(), &value_validity));
+  ARROW_RETURN_NOT_OK(decomposer.Decompose(*expr->root(), &value_validity));
 
   auto& used_functions = decomposer.UsedFunctions();
-  functions_in_exprs_.insert(
-      used_functions.begin(),
-      used_functions.end());
+  functions_in_exprs_.insert(used_functions.begin(), used_functions.end());
 
   return value_validity;
 }
 
-Status LLVMGenerator::Add(const ExpressionPtr expr, ValueValidityPairPtr value_validity, const FieldDescriptorPtr output) {
+Status LLVMGenerator::Add(const ExpressionPtr expr, ValueValidityPairPtr value_validity,
+                          const FieldDescriptorPtr output) {
   int idx = static_cast<int>(compiled_exprs_.size());
   // Generate the IR function for the decomposed expression.
   auto compiled_expr = std::make_unique<CompiledExpr>(value_validity, output);

cpp/src/gandiva/llvm_generator.h

@@ -194,7 +194,7 @@ class GANDIVA_EXPORT LLVMGenerator {
 
   // Generate the code for one expression for default mode, with the output of
   // the expression going to 'output'.
-  Status Add(const ExpressionPtr expr, ValueValidityPairPtr value_validity, 
+  Status Add(const ExpressionPtr expr, ValueValidityPairPtr value_validity,
              const FieldDescriptorPtr output);
 
   /// Generate code to load the vector at specified index in the 'arg_addrs' array.