GH-40024: [C++][Gandiva] Selectively register external C functions based on expression usage

cpp/src/gandiva/engine.cc

@@ -346,11 +346,21 @@ Engine::Engine(const std::shared_ptr<Configuration>& conf,
 
 Engine::~Engine() {}
 
-Status Engine::Init() {
+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();
 }
 
@@ -599,6 +609,10 @@ 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) {
+ auto* existing = module()->getFunction(name);
+ if (existing != nullptr) {
+ 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

@@ -93,13 +93,17 @@ class GANDIVA_EXPORT Engine {
  // Create a global string as a pointer with "i8*" type.
  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);
 
  // Post construction init. This _must_ be called after the constructor.
- Status Init();
 
  static void InitOnce();
 
@@ -123,11 +127,31 @@ class GANDIVA_EXPORT Engine {
  LLVMTypes types_;
 
  std::vector<std::string> functions_to_compile_;
+ 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",
+ };
 
  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/expr_decomposer.cc

@@ -70,6 +70,7 @@ Status ExprDecomposer::Visit(const FunctionNode& in_node) {
  const NativeFunction* native_function = registry_.LookupSignature(signature);
  DCHECK(native_function) << "Missing Signature " << signature.ToString();
 
+ used_functions_.emplace(native_function->pc_name());
  // decompose the children.
  std::vector<ValueValidityPairPtr> args;
  for (auto& child : node.children()) {

cpp/src/gandiva/expr_decomposer.h

@@ -21,6 +21,7 @@
 #include <memory>
 #include <stack>
 #include <string>
+#include <unordered_set>
 #include <utility>
 
 #include "gandiva/arrow.h"
@@ -49,6 +50,10 @@ class GANDIVA_EXPORT ExprDecomposer : public NodeVisitor {
  return status;
  }
 
+ [[nodiscard]] const std::unordered_set<std::string>& UsedFunctions() const {
+ return used_functions_;
+ }
+
  private:
  ARROW_DISALLOW_COPY_AND_ASSIGN(ExprDecomposer);
 
@@ -125,6 +130,7 @@ class GANDIVA_EXPORT ExprDecomposer : public NodeVisitor {
  Annotator& annotator_;
  std::stack<std::unique_ptr<IfStackEntry>> if_entries_stack_;
  ValueValidityPairPtr result_;
+ std::unordered_set<std::string> used_functions_;
  bool nested_if_else_;
 };
 

cpp/src/gandiva/expr_decomposer_test.cc

@@ -19,11 +19,13 @@
 
 #include <gtest/gtest.h>
 
+#include "arrow/testing/gtest_util.h"
 #include "gandiva/annotator.h"
 #include "gandiva/dex.h"
 #include "gandiva/function_registry.h"
 #include "gandiva/gandiva_aliases.h"
 #include "gandiva/node.h"
+#include "gandiva/tree_expr_builder.h"
 
 namespace gandiva {
 
@@ -405,4 +407,34 @@ TEST_F(TestExprDecomposer, TestComplexIfCondition) {
  EXPECT_TRUE(decomposer.if_entries_stack_.empty());
 }
 
+TEST_F(TestExprDecomposer, TestGetUsedFunctionsInExpr) {
+ Annotator annotator;
+ ExprDecomposer decomposer(*registry_, annotator);
+ auto field0 = field("f0", int32());
+ auto f0 = std::make_shared<FieldNode>(field0);
+ auto is_not_null_func = TreeExprBuilder::MakeFunction("isnotnull", {f0}, boolean());
+ ValueValidityPairPtr value_validity;
+ ASSERT_OK(decomposer.Decompose(*is_not_null_func, &value_validity));
+ auto used_functions = decomposer.UsedFunctions();
+ ASSERT_EQ(used_functions.size(), 1);
+ ASSERT_EQ(used_functions.find("isnotnull_int32") != used_functions.end(), true);
+}
+
+TEST_F(TestExprDecomposer, TestGetMultipleUsedFunctionsInExpr) {
+ Annotator annotator;
+ ExprDecomposer decomposer(*registry_, annotator);
+ auto field0 = field("f0", int32());
+ auto field1 = field("f1", int32());
+ auto f0 = std::make_shared<FieldNode>(field0);
+ auto f1 = std::make_shared<FieldNode>(field1);
+ auto add_func = TreeExprBuilder::MakeFunction("add", {f0, f1}, int32());
+ auto is_not_null_func =
+ TreeExprBuilder::MakeFunction("isnotnull", {add_func}, boolean());
+ ValueValidityPairPtr value_validity;
+ ASSERT_OK(decomposer.Decompose(*is_not_null_func, &value_validity));
+ auto used_functions = decomposer.UsedFunctions();
+ ASSERT_EQ(used_functions.size(), 2);
+ ASSERT_EQ(used_functions.find("add_int32_int32") != used_functions.end(), true);
+ ASSERT_EQ(used_functions.find("isnotnull_int32") != used_functions.end(), true);
+}
 } // namespace gandiva

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,12 +67,21 @@ Status LLVMGenerator::SetLLVMObjectCache(GandivaObjectCache& object_cache) {
  return engine_->SetLLVMObjectCache(object_cache);
 }
 
-Status LLVMGenerator::Add(const ExpressionPtr expr, const FieldDescriptorPtr output) {
- int idx = static_cast<int>(compiled_exprs_.size());
- // decompose the expression to separate out value and validities.
+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));
+
+ auto& used_functions = decomposer.UsedFunctions();
+ 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) {
+ 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);
  std::string fn_name = "expr_" + std::to_string(idx) + "_" +
@@ -93,9 +102,19 @@ Status LLVMGenerator::Add(const ExpressionPtr expr, const FieldDescriptorPtr out
 Status LLVMGenerator::Build(const ExpressionVector& exprs, SelectionVector::Mode mode) {
  selection_vector_mode_ = mode;
 
+ std::vector<ValueValidityPairPtr> expr_value_validities;
  for (auto& expr : exprs) {
+ ARROW_ASSIGN_OR_RAISE(auto value_validity, Decompose(expr));
+ expr_value_validities.push_back(value_validity);
+ }
+
+ ARROW_RETURN_NOT_OK(engine_->Init(std::move(functions_in_exprs_)));
+
+ for (size_t i = 0; i < exprs.size(); ++i) {
+ const auto& expr = exprs[i];
  auto output = annotator_.AddOutputFieldDescriptor(expr->result());
- ARROW_RETURN_NOT_OK(Add(expr, output));
+ auto value_validity = expr_value_validities[i];
+ ARROW_RETURN_NOT_OK(Add(expr, std::move(value_validity), output));
  }
 
  // Compile and inject into the process' memory the generated function.

cpp/src/gandiva/llvm_generator.h

@@ -190,9 +190,12 @@ class GANDIVA_EXPORT LLVMGenerator {
  bool has_arena_allocs_;
  };
 
+ arrow::Result<ValueValidityPairPtr> Decompose(const ExpressionPtr& expr);
+
  // Generate the code for one expression for default mode, with the output of
  // the expression going to 'output'.
- Status Add(const ExpressionPtr expr, const FieldDescriptorPtr output);
+ 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.
  llvm::Value* LoadVectorAtIndex(llvm::Value* arg_addrs, llvm::Type* type, int idx,
@@ -269,6 +272,7 @@ class GANDIVA_EXPORT LLVMGenerator {
  // used for debug
  bool enable_ir_traces_;
  std::vector<std::string> trace_strings_;
+ std::unordered_set<std::string> functions_in_exprs_;
 };
 
 } // namespace gandiva

cpp/src/gandiva/tests/micro_benchmarks.cc

@@ -450,6 +450,29 @@ static void TimedTestExprCompilation(benchmark::State& state) {
  }
 }
 
+static void TimedTestNonBitcodeExprCompilation(benchmark::State& state, bool use_cache) {
+ int32_t iteration = 0;
+ for (auto _ : state) {
+ // schema for input fields
+ double literal_value = use_cache ? 1.0 : static_cast<double>(iteration);
+ auto seed = TreeExprBuilder::MakeLiteral(literal_value);
+ auto schema = arrow::schema({});
+
+ // output field
+ auto field_sin = field("c1", float64());
+
+ // seed is different for each iteration so that cache won't be hit
+ auto sin_func = TreeExprBuilder::MakeFunction("sin", {seed}, float64());
+
+ auto expr_0 = TreeExprBuilder::MakeExpression(sin_func, field_sin);
+
+ std::shared_ptr<Projector> projector;
+ ASSERT_OK(Projector::Make(schema, {expr_0}, TestConfiguration(), &projector));
+
+ ++iteration;
+ }
+}
+
 static void DecimalAdd2Fast(benchmark::State& state) {
  // use lesser precision to test the fast-path
  DoDecimalAdd2(state, DecimalTypeUtil::kMaxPrecision - 6, 18);
@@ -490,6 +513,16 @@ static void DecimalAdd3Large(benchmark::State& state) {
  DoDecimalAdd3(state, DecimalTypeUtil::kMaxPrecision, 18, true);
 }
 
+static void TimedTestNonBitcodeExprCompilationNoCache(benchmark::State& state) {
+ TimedTestNonBitcodeExprCompilation(state, false);
+}
+
+static void TimedTestNonBitcodeExprCompilationWithCache(benchmark::State& state) {
+ TimedTestNonBitcodeExprCompilation(state, true);
+}
+
+BENCHMARK(TimedTestNonBitcodeExprCompilationNoCache)->Unit(benchmark::kMicrosecond);
+BENCHMARK(TimedTestNonBitcodeExprCompilationWithCache)->Unit(benchmark::kMicrosecond);
 BENCHMARK(TimedTestExprCompilation)->Unit(benchmark::kMicrosecond);
 BENCHMARK(TimedTestAdd3)->Unit(benchmark::kMicrosecond);
 BENCHMARK(TimedTestBigNested)->Unit(benchmark::kMicrosecond);