GH-55281: [C++][Gandiva] Release arena memory per row for any function with context

cpp/src/gandiva/execution_context.h

@@ -41,6 +41,17 @@ class ExecutionContext {
 
   SimpleArena* arena() { return &arena_; }
 
+  // Byte offset of the arena's "used" flag from the start of a context. Generated code
+  // reads the flag through the context pointer it already has.
+  static int64_t arena_used_offset() {
+    static const int64_t offset = [] {
+      ExecutionContext context;
+      return reinterpret_cast<char*>(context.arena()->used_ptr()) -
+             reinterpret_cast<char*>(&context);
+    }();
+    return offset;
+  }
+
   void Reset() {
     error_msg_.clear();
     arena_.Reset();

cpp/src/gandiva/llvm_generator.cc

@@ -25,6 +25,7 @@
 #include "gandiva/bitmap_accumulator.h"
 #include "gandiva/decimal_ir.h"
 #include "gandiva/dex.h"
+#include "gandiva/execution_context.h"
 #include "gandiva/expr_decomposer.h"
 #include "gandiva/expression.h"
 #include "gandiva/llvm_types.h"
@@ -427,10 +428,26 @@ Status LLVMGenerator::CodeGenExprValue(DexPtr value_expr, int buffer_count,
   if (visitor.has_arena_allocs()) {
     // Reset allocations to avoid excessive memory usage. Once the result is copied to
     // the output vector (store instruction above), any memory allocations in this
-    // iteration of the loop are no longer needed.
-    std::vector<llvm::Value*> reset_args;
-    reset_args.push_back(arg_context_ptr);
-    AddFunctionCall("gdv_fn_context_arena_reset", types()->void_type(), reset_args);
+    // iteration of the loop are no longer needed. Skip the call when the row did not
+    // allocate, which keeps cheap functions that only report errors fast.
+    llvm::BasicBlock* arena_reset =
+        llvm::BasicBlock::Create(*context(), "arena_reset", fn, loop_exit);
+    llvm::BasicBlock* loop_latch =
+        llvm::BasicBlock::Create(*context(), "loop_latch", fn, loop_exit);
+    llvm::Value* used_addr = builder->CreateAdd(
+        arg_context_ptr, types()->i64_constant(ExecutionContext::arena_used_offset()));
+    llvm::Value* arena_used = builder->CreateLoad(
+        types()->i8_type(), builder->CreateIntToPtr(used_addr, types()->i8_ptr_type()),
+        "arena_used");
+    builder->CreateCondBr(
+        builder->CreateICmpNE(arena_used, types()->i8_constant(0), "arena_used?"),
+        arena_reset, loop_latch);
+    builder->SetInsertPoint(arena_reset);
+    AddFunctionCall("gdv_fn_context_arena_reset", types()->void_type(),
+                    {arg_context_ptr});
+    builder->CreateBr(loop_latch);
+    builder->SetInsertPoint(loop_latch);
+    loop_body_tail = loop_latch;
   }
 
   // check loop_var
@@ -1263,6 +1280,12 @@ LValuePtr LLVMGenerator::Visitor::BuildFunctionCall(const NativeFunction* func,
   auto llvm_return_type = types->IRType(arrow_return_type_id);
   DecimalIR decimalIR(generator_->engine_.get());
 
+  // Only functions that get the execution context can allocate from its arena (directly
+  // or through a function holder), so those need the arena released after each row.
+  if (func->NeedsContext()) {
+    has_arena_allocs_ = true;
+  }
+
   if (arrow_return_type_id == arrow::Type::DECIMAL) {
     // For decimal fns, the output precision/scale are passed along as parameters.
     //
@@ -1293,7 +1316,6 @@ LValuePtr LLVMGenerator::Visitor::BuildFunctionCall(const NativeFunction* func,
       result_len_ptr = new llvm::AllocaInst(generator_->types()->i32_type(), 0,
                                             "result_len", entry_block_);
       params->push_back(result_len_ptr);
-      has_arena_allocs_ = true;
     }
 
     // Make the function call

cpp/src/gandiva/simple_arena.h

@@ -57,6 +57,10 @@ class SimpleArena {
   // total bytes available for allocations.
   int64_t avail_bytes() { return avail_bytes_; }
 
+  // Non-zero if Allocate() was called since the last Reset(). Generated code reads it to
+  // skip the Reset() call when nothing was allocated.
+  uint8_t* used_ptr() { return &used_; }
+
  private:
   struct Chunk {
     Chunk(uint8_t* buf, int64_t size) : buf_(buf), size_(size) {}
@@ -88,18 +92,23 @@ class SimpleArena {
 
   // List of allocated chunks.
   std::vector<Chunk> chunks_;
+
+  // Non-zero if Allocate() was called since the last Reset().
+  uint8_t used_;
 };
 
 inline SimpleArena::SimpleArena(arrow::MemoryPool* pool, int64_t min_chunk_size)
     : pool_(pool),
       min_chunk_size_(min_chunk_size),
       total_bytes_(0),
       avail_bytes_(0),
-      avail_buf_(NULL) {}
+      avail_buf_(NULL),
+      used_(0) {}
 
 inline SimpleArena::~SimpleArena() { ReleaseChunks(false /*retain_first*/); }
 
 inline uint8_t* SimpleArena::Allocate(int64_t size) {
+  used_ = 1;
   if (avail_bytes_ < size) {
     auto status = AllocateChunk(std::max(size, min_chunk_size_));
     if (!status.ok()) {
@@ -131,6 +140,7 @@ inline arrow::Status SimpleArena::AllocateChunk(int64_t size) {
 // By retaining the first chunk, the number of malloc calls are reduced to one per batch,
 // instead of one per record.
 inline void SimpleArena::Reset() {
+  used_ = 0;
   if (chunks_.size() == 0) {
     // if there are no chunks, nothing to do.
     return;

cpp/src/gandiva/tests/projector_test.cc

@@ -1276,6 +1276,96 @@ TEST_F(TestProjector, TestLevenshtein) {
   EXPECT_ARROW_ARRAY_EQUALS(exp_lev, outputs.at(0));
 }
 
+// Per-row arena scratch space is released after each row only if the expression contains
+// a function that returns a string. levenshtein returns an int32 but takes
+// (shorter_len + 1) * 2 ints of arena scratch per row, so on its own it holds that
+// memory until the end of the batch.
+TEST_F(TestProjector, TestLevenshteinArenaGrowth) {
+  constexpr int kNumRecords = 1 << 18;
+  constexpr int kStringLen = 16;
+  // About 35 MiB of scratch for the whole batch if it is not released row by row.
+  constexpr int64_t kScratchBytes =
+      static_cast<int64_t>(kNumRecords) * (kStringLen + 1) * 2 * sizeof(int);
+  // Outputs and one arena chunk need about 1 MiB when it is released, so this leaves a
+  // wide margin on both sides.
+  constexpr int64_t kMaxBatchBytes = kScratchBytes / 4;
+
+  auto f0 = field("f0", arrow::utf8());
+  auto f1 = field("f1", arrow::utf8());
+  auto schema = arrow::schema({f0, f1});
+  auto field_lev = field("levenshtein", arrow::int32());
+
+  std::string left(kStringLen, 'a');
+  std::string right = left;
+  right[0] = 'b';
+  std::vector<std::string> lefts(kNumRecords, left);
+  std::vector<std::string> rights(kNumRecords, right);
+  std::vector<bool> validity(kNumRecords, true);
+  std::vector<int32_t> expected_values(kNumRecords, 1);
+  auto in_batch = arrow::RecordBatch::Make(
+      schema, kNumRecords,
+      {MakeArrowArrayUtf8(lefts, validity), MakeArrowArrayUtf8(rights, validity)});
+  auto expected = MakeArrowArrayInt32(expected_values, validity);
+
+  auto expression = TreeExprBuilder::MakeExpression("levenshtein", {f0, f1}, field_lev);
+  std::shared_ptr<Projector> projector;
+  ASSERT_OK(Projector::Make(schema, {expression}, TestConfiguration(), &projector));
+
+  // The execution context's arena draws from the default pool, whichever pool is
+  // passed to Evaluate, and returns everything by the end of the call. Measure the
+  // bytes requested from the default pool during the call, not the live bytes.
+  auto* default_pool = arrow::default_memory_pool();
+  const int64_t before = default_pool->total_bytes_allocated();
+  arrow::ArrayVector outputs;
+  ASSERT_OK(projector->Evaluate(*in_batch, pool_, &outputs));
+  const int64_t allocated = default_pool->total_bytes_allocated() - before;
+
+  ASSERT_EQ(outputs.size(), 1);
+  EXPECT_ARROW_ARRAY_EQUALS(expected, outputs[0]);
+  EXPECT_LT(allocated, kMaxBatchBytes)
+      << "allocated " << allocated / (1 << 20) << " MiB while evaluating " << kNumRecords
+      << " rows; the arena is not released per row";
+}
+
+// upper() takes its result from the arena, but the expression returns an int32. The arena
+// must still be released after each row.
+TEST_F(TestProjector, TestNestedStringFunctionArenaGrowth) {
+  constexpr int kNumRecords = 1 << 18;
+  constexpr int kStringLen = 64;
+  // 16 MiB of upper() results for the whole batch if they are not released row by row.
+  constexpr int64_t kScratchBytes = static_cast<int64_t>(kNumRecords) * kStringLen;
+  // Outputs and one arena chunk need about 1 MiB when it is released.
+  constexpr int64_t kMaxBatchBytes = kScratchBytes / 4;
+
+  auto in_field = field("in", arrow::utf8());
+  auto schema = arrow::schema({in_field});
+  auto upper = TreeExprBuilder::MakeFunction(
+      "upper", {TreeExprBuilder::MakeField(in_field)}, arrow::utf8());
+  auto length = TreeExprBuilder::MakeFunction("length", {upper}, arrow::int32());
+  auto expr = TreeExprBuilder::MakeExpression(length, field("out", arrow::int32()));
+
+  std::shared_ptr<Projector> projector;
+  ASSERT_OK(Projector::Make(schema, {expr}, TestConfiguration(), &projector));
+
+  std::vector<std::string> values(kNumRecords, std::string(kStringLen, 'a'));
+  std::vector<bool> validity(kNumRecords, true);
+  auto in_batch = arrow::RecordBatch::Make(schema, kNumRecords,
+                                           {MakeArrowArrayUtf8(values, validity)});
+  auto expected =
+      MakeArrowArrayInt32(std::vector<int32_t>(kNumRecords, kStringLen), validity);
+
+  auto* default_pool = arrow::default_memory_pool();
+  const int64_t before = default_pool->total_bytes_allocated();
+  arrow::ArrayVector outputs;
+  ASSERT_OK(projector->Evaluate(*in_batch, pool_, &outputs));
+  const int64_t allocated = default_pool->total_bytes_allocated() - before;
+
+  EXPECT_ARROW_ARRAY_EQUALS(expected, outputs.at(0));
+  EXPECT_LT(allocated, kMaxBatchBytes)
+      << "allocated " << allocated / (1 << 20) << " MiB while evaluating " << kNumRecords
+      << " rows; the arena is not released per row";
+}
+
 TEST_F(TestProjector, TestQuote) {
   // schema for input fields
   auto field0 = field("f0", arrow::utf8());
@@ -3761,6 +3851,45 @@ TEST_F(TestProjector, TestExtendedCFunctionThatNeedsContext) {
   EXPECT_ARROW_ARRAY_EQUALS(out, outs.at(0));
 }
 
+// scratch_length takes arena scratch memory but returns an int32. The arena must still be
+// released after each row.
+TEST_F(TestProjector, TestExtendedCFunctionArenaGrowth) {
+  constexpr int kNumRecords = 1 << 18;
+  constexpr int kStringLen = 64;
+  // 16 MiB of scratch for the whole batch if it is not released row by row.
+  constexpr int64_t kScratchBytes = static_cast<int64_t>(kNumRecords) * kStringLen;
+  // Outputs and one arena chunk need about 1 MiB when it is released.
+  constexpr int64_t kMaxBatchBytes = kScratchBytes / 4;
+
+  auto in_field = field("in", arrow::utf8());
+  auto schema = arrow::schema({in_field});
+  auto out_field = field("out", arrow::int32());
+  auto expr = TreeExprBuilder::MakeExpression("scratch_length", {in_field}, out_field);
+
+  std::shared_ptr<Projector> projector;
+  auto external_registry = std::make_shared<FunctionRegistry>();
+  auto config = TestConfigWithArenaScratchFunction(std::move(external_registry));
+  ASSERT_OK(Projector::Make(schema, {expr}, config, &projector));
+
+  std::vector<std::string> values(kNumRecords, std::string(kStringLen, 'a'));
+  std::vector<bool> validity(kNumRecords, true);
+  auto in_batch = arrow::RecordBatch::Make(schema, kNumRecords,
+                                           {MakeArrowArrayUtf8(values, validity)});
+  auto expected =
+      MakeArrowArrayInt32(std::vector<int32_t>(kNumRecords, kStringLen), validity);
+
+  auto* default_pool = arrow::default_memory_pool();
+  const int64_t before = default_pool->total_bytes_allocated();
+  arrow::ArrayVector outputs;
+  ASSERT_OK(projector->Evaluate(*in_batch, pool_, &outputs));
+  const int64_t allocated = default_pool->total_bytes_allocated() - before;
+
+  EXPECT_ARROW_ARRAY_EQUALS(expected, outputs.at(0));
+  EXPECT_LT(allocated, kMaxBatchBytes)
+      << "allocated " << allocated / (1 << 20) << " MiB while evaluating " << kNumRecords
+      << " rows; the arena is not released per row";
+}
+
 TEST_F(TestProjector, TestRandomNoArgs) {
   // Test random() with no arguments - returns double in [0, 1)
   auto dummy_field = field("dummy", arrow::int32());

cpp/src/gandiva/tests/test_util.cc

@@ -76,6 +76,13 @@ static NativeFunction GetTestFunctionWithContext() {
   return multiply_by_two_formula;
 }
 
+static NativeFunction GetTestFunctionWithArenaScratch() {
+  NativeFunction scratch_length("scratch_length", {}, {arrow::utf8()}, arrow::int32(),
+                                ResultNullableType::kResultNullIfNull,
+                                "scratch_length_utf8", NativeFunction::kNeedsContext);
+  return scratch_length;
+}
+
 static std::shared_ptr<Configuration> BuildConfigurationWithRegistry(
     std::shared_ptr<FunctionRegistry> registry,
     const std::function<arrow::Status(std::shared_ptr<FunctionRegistry>)>&
@@ -146,6 +153,18 @@ static const char* multiply_by_two_formula(int64_t ctx, const char* value,
   memcpy(out, result.c_str(), *out_len);
   return out;
 }
+
+// copies the value into arena scratch memory and returns its length, so it allocates
+// from the arena without returning a string
+static int32_t scratch_length(int64_t ctx, const char* value, int32_t value_len) {
+  auto scratch = reinterpret_cast<char*>(gdv_fn_context_arena_malloc(ctx, value_len));
+  if (scratch == nullptr) {
+    gdv_fn_context_set_error_msg(ctx, "Could not allocate scratch memory");
+    return 0;
+  }
+  memcpy(scratch, value, value_len);
+  return value_len;
+}
 }
 
 std::shared_ptr<Configuration> TestConfigWithCFunction(
@@ -172,4 +191,12 @@ std::shared_ptr<Configuration> TestConfigWithContextFunction(
                          reinterpret_cast<void*>(multiply_by_two_formula));
   });
 }
+
+std::shared_ptr<Configuration> TestConfigWithArenaScratchFunction(
+    std::shared_ptr<FunctionRegistry> registry) {
+  return BuildConfigurationWithRegistry(std::move(registry), [](auto reg) {
+    return reg->Register(GetTestFunctionWithArenaScratch(),
+                         reinterpret_cast<void*>(scratch_length));
+  });
+}
 }  // namespace gandiva

cpp/src/gandiva/tests/test_util.h

@@ -122,6 +122,12 @@ std::shared_ptr<Configuration> TestConfigWithHolderFunction(
 std::shared_ptr<Configuration> TestConfigWithContextFunction(
     std::shared_ptr<FunctionRegistry> registry);
 
+// helper function to create a Configuration with an external function registered
+// to the given function registry, and the external function takes arena scratch memory
+// but returns an int32
+std::shared_ptr<Configuration> TestConfigWithArenaScratchFunction(
+    std::shared_ptr<FunctionRegistry> registry);
+
 std::string GetTestFunctionLLVMIRPath();
 NativeFunction GetTestExternalFunction();