From 4671b2ad8722ca2cbdded67736446a7ce8922df6 Mon Sep 17 00:00:00 2001 From: Dmitry Chirkov Date: Mon, 26 Oct 2026 22:32:47 -0700 Subject: [PATCH 1/4] GH-55281: [C++][Gandiva] Add regression test for levenshtein arena growth levenshtein returns int32 but uses arena memory, so the per-row arena release is never emitted for it. This test fails until that is fixed. --- cpp/src/gandiva/tests/projector_test.cc | 51 +++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/cpp/src/gandiva/tests/projector_test.cc b/cpp/src/gandiva/tests/projector_test.cc index 6e9450e1c650..03ad72522026 100644 --- a/cpp/src/gandiva/tests/projector_test.cc +++ b/cpp/src/gandiva/tests/projector_test.cc @@ -1276,6 +1276,57 @@ 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(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 lefts(kNumRecords, left); + std::vector rights(kNumRecords, right); + std::vector validity(kNumRecords, true); + std::vector 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; + 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"; +} + TEST_F(TestProjector, TestQuote) { // schema for input fields auto field0 = field("f0", arrow::utf8()); From 95d67a9e74b1f88fe30d808f08b0f0b2918abcb0 Mon Sep 17 00:00:00 2001 From: Dmitry Chirkov Date: Mon, 26 Oct 2026 22:30:48 -0700 Subject: [PATCH 2/4] GH-55281: [C++][Gandiva] Release arena memory per row for levenshtein Add a kUsesArena flag, set it for levenshtein, and release the arena per row for flagged functions. --- cpp/src/gandiva/function_registry_string.cc | 3 ++- cpp/src/gandiva/llvm_generator.cc | 6 ++++++ cpp/src/gandiva/native_function.h | 3 +++ 3 files changed, 11 insertions(+), 1 deletion(-) diff --git a/cpp/src/gandiva/function_registry_string.cc b/cpp/src/gandiva/function_registry_string.cc index 35938cb82908..ab816b2fe33d 100644 --- a/cpp/src/gandiva/function_registry_string.cc +++ b/cpp/src/gandiva/function_registry_string.cc @@ -480,7 +480,8 @@ std::vector GetStringFunctionRegistry() { NativeFunction("levenshtein", {}, DataTypeVector{utf8(), utf8()}, int32(), kResultNullIfNull, "levenshtein", - NativeFunction::kNeedsContext | NativeFunction::kCanReturnErrors), + NativeFunction::kNeedsContext | NativeFunction::kCanReturnErrors | + NativeFunction::kUsesArena), NativeFunction("to_hex", {"hex"}, DataTypeVector{binary()}, utf8(), kResultNullIfNull, "to_hex_binary", NativeFunction::kNeedsContext), diff --git a/cpp/src/gandiva/llvm_generator.cc b/cpp/src/gandiva/llvm_generator.cc index 7aa91c4dd6b4..2e2d5c46eb01 100644 --- a/cpp/src/gandiva/llvm_generator.cc +++ b/cpp/src/gandiva/llvm_generator.cc @@ -1263,6 +1263,12 @@ LValuePtr LLVMGenerator::Visitor::BuildFunctionCall(const NativeFunction* func, auto llvm_return_type = types->IRType(arrow_return_type_id); DecimalIR decimalIR(generator_->engine_.get()); + // The arena is released per row, so a function that allocates from it needs that + // even when it does not return a string. + if (func->UsesArena()) { + has_arena_allocs_ = true; + } + if (arrow_return_type_id == arrow::Type::DECIMAL) { // For decimal fns, the output precision/scale are passed along as parameters. // diff --git a/cpp/src/gandiva/native_function.h b/cpp/src/gandiva/native_function.h index c20de3dbdd54..25f86b89c710 100644 --- a/cpp/src/gandiva/native_function.h +++ b/cpp/src/gandiva/native_function.h @@ -44,6 +44,8 @@ class GANDIVA_EXPORT NativeFunction { static constexpr int32_t kNeedsContext = (1 << 1); static constexpr int32_t kNeedsFunctionHolder = (1 << 2); static constexpr int32_t kCanReturnErrors = (1 << 3); + // Allocates from the context arena without returning a string or binary. + static constexpr int32_t kUsesArena = (1 << 4); const std::vector& signatures() const { return signatures_; } std::string pc_name() const { return pc_name_; } @@ -52,6 +54,7 @@ class GANDIVA_EXPORT NativeFunction { bool NeedsContext() const { return (flags_ & kNeedsContext) != 0; } bool NeedsFunctionHolder() const { return (flags_ & kNeedsFunctionHolder) != 0; } bool CanReturnErrors() const { return (flags_ & kCanReturnErrors) != 0; } + bool UsesArena() const { return (flags_ & kUsesArena) != 0; } NativeFunction(const std::string& base_name, const std::vector& aliases, const DataTypeVector& param_types, const DataTypePtr& ret_type, From c8afba88c1ef7e1bc384d6bdbb96fbd1418e3169 Mon Sep 17 00:00:00 2001 From: Dmitry Chirkov Date: Tue, 27 Oct 2026 00:28:51 -0700 Subject: [PATCH 3/4] GH-55281: [C++][Gandiva] Add arena tests for external and nested string functions An external function that takes the context and arena scratch but returns int32 still keeps its memory until the end of the batch. This test fails until that is fixed. --- cpp/src/gandiva/tests/projector_test.cc | 78 +++++++++++++++++++++++++ cpp/src/gandiva/tests/test_util.cc | 27 +++++++++ cpp/src/gandiva/tests/test_util.h | 6 ++ 3 files changed, 111 insertions(+) diff --git a/cpp/src/gandiva/tests/projector_test.cc b/cpp/src/gandiva/tests/projector_test.cc index 03ad72522026..afb0d503b183 100644 --- a/cpp/src/gandiva/tests/projector_test.cc +++ b/cpp/src/gandiva/tests/projector_test.cc @@ -1327,6 +1327,45 @@ TEST_F(TestProjector, TestLevenshteinArenaGrowth) { << " 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(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; + ASSERT_OK(Projector::Make(schema, {expr}, TestConfiguration(), &projector)); + + std::vector values(kNumRecords, std::string(kStringLen, 'a')); + std::vector validity(kNumRecords, true); + auto in_batch = arrow::RecordBatch::Make(schema, kNumRecords, + {MakeArrowArrayUtf8(values, validity)}); + auto expected = + MakeArrowArrayInt32(std::vector(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()); @@ -3812,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(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; + auto external_registry = std::make_shared(); + auto config = TestConfigWithArenaScratchFunction(std::move(external_registry)); + ASSERT_OK(Projector::Make(schema, {expr}, config, &projector)); + + std::vector values(kNumRecords, std::string(kStringLen, 'a')); + std::vector validity(kNumRecords, true); + auto in_batch = arrow::RecordBatch::Make(schema, kNumRecords, + {MakeArrowArrayUtf8(values, validity)}); + auto expected = + MakeArrowArrayInt32(std::vector(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()); diff --git a/cpp/src/gandiva/tests/test_util.cc b/cpp/src/gandiva/tests/test_util.cc index 3d0ab51752e1..34d39add6e4c 100644 --- a/cpp/src/gandiva/tests/test_util.cc +++ b/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 BuildConfigurationWithRegistry( std::shared_ptr registry, const std::function)>& @@ -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(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 TestConfigWithCFunction( @@ -172,4 +191,12 @@ std::shared_ptr TestConfigWithContextFunction( reinterpret_cast(multiply_by_two_formula)); }); } + +std::shared_ptr TestConfigWithArenaScratchFunction( + std::shared_ptr registry) { + return BuildConfigurationWithRegistry(std::move(registry), [](auto reg) { + return reg->Register(GetTestFunctionWithArenaScratch(), + reinterpret_cast(scratch_length)); + }); +} } // namespace gandiva diff --git a/cpp/src/gandiva/tests/test_util.h b/cpp/src/gandiva/tests/test_util.h index d8181fe67516..beb92e626fd0 100644 --- a/cpp/src/gandiva/tests/test_util.h +++ b/cpp/src/gandiva/tests/test_util.h @@ -122,6 +122,12 @@ std::shared_ptr TestConfigWithHolderFunction( std::shared_ptr TestConfigWithContextFunction( std::shared_ptr 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 TestConfigWithArenaScratchFunction( + std::shared_ptr registry); + std::string GetTestFunctionLLVMIRPath(); NativeFunction GetTestExternalFunction(); From e2a8f11529361440123ff16823823675306b6f0b Mon Sep 17 00:00:00 2001 From: Dmitry Chirkov Date: Tue, 27 Oct 2026 00:30:00 -0700 Subject: [PATCH 4/4] GH-55281: [C++][Gandiva] Release arena memory per row for any function with context Use NeedsContext() instead of the kUsesArena flag, and skip the reset call when the row did not allocate. --- cpp/src/gandiva/execution_context.h | 11 +++++++ cpp/src/gandiva/function_registry_string.cc | 3 +- cpp/src/gandiva/llvm_generator.cc | 32 +++++++++++++++------ cpp/src/gandiva/native_function.h | 3 -- cpp/src/gandiva/simple_arena.h | 12 +++++++- 5 files changed, 47 insertions(+), 14 deletions(-) diff --git a/cpp/src/gandiva/execution_context.h b/cpp/src/gandiva/execution_context.h index efa546874456..5e3d6f3fe3b9 100644 --- a/cpp/src/gandiva/execution_context.h +++ b/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(context.arena()->used_ptr()) - + reinterpret_cast(&context); + }(); + return offset; + } + void Reset() { error_msg_.clear(); arena_.Reset(); diff --git a/cpp/src/gandiva/function_registry_string.cc b/cpp/src/gandiva/function_registry_string.cc index ab816b2fe33d..35938cb82908 100644 --- a/cpp/src/gandiva/function_registry_string.cc +++ b/cpp/src/gandiva/function_registry_string.cc @@ -480,8 +480,7 @@ std::vector GetStringFunctionRegistry() { NativeFunction("levenshtein", {}, DataTypeVector{utf8(), utf8()}, int32(), kResultNullIfNull, "levenshtein", - NativeFunction::kNeedsContext | NativeFunction::kCanReturnErrors | - NativeFunction::kUsesArena), + NativeFunction::kNeedsContext | NativeFunction::kCanReturnErrors), NativeFunction("to_hex", {"hex"}, DataTypeVector{binary()}, utf8(), kResultNullIfNull, "to_hex_binary", NativeFunction::kNeedsContext), diff --git a/cpp/src/gandiva/llvm_generator.cc b/cpp/src/gandiva/llvm_generator.cc index 2e2d5c46eb01..f787cf8e4872 100644 --- a/cpp/src/gandiva/llvm_generator.cc +++ b/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 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,9 +1280,9 @@ LValuePtr LLVMGenerator::Visitor::BuildFunctionCall(const NativeFunction* func, auto llvm_return_type = types->IRType(arrow_return_type_id); DecimalIR decimalIR(generator_->engine_.get()); - // The arena is released per row, so a function that allocates from it needs that - // even when it does not return a string. - if (func->UsesArena()) { + // 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; } @@ -1299,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 diff --git a/cpp/src/gandiva/native_function.h b/cpp/src/gandiva/native_function.h index 25f86b89c710..c20de3dbdd54 100644 --- a/cpp/src/gandiva/native_function.h +++ b/cpp/src/gandiva/native_function.h @@ -44,8 +44,6 @@ class GANDIVA_EXPORT NativeFunction { static constexpr int32_t kNeedsContext = (1 << 1); static constexpr int32_t kNeedsFunctionHolder = (1 << 2); static constexpr int32_t kCanReturnErrors = (1 << 3); - // Allocates from the context arena without returning a string or binary. - static constexpr int32_t kUsesArena = (1 << 4); const std::vector& signatures() const { return signatures_; } std::string pc_name() const { return pc_name_; } @@ -54,7 +52,6 @@ class GANDIVA_EXPORT NativeFunction { bool NeedsContext() const { return (flags_ & kNeedsContext) != 0; } bool NeedsFunctionHolder() const { return (flags_ & kNeedsFunctionHolder) != 0; } bool CanReturnErrors() const { return (flags_ & kCanReturnErrors) != 0; } - bool UsesArena() const { return (flags_ & kUsesArena) != 0; } NativeFunction(const std::string& base_name, const std::vector& aliases, const DataTypeVector& param_types, const DataTypePtr& ret_type, diff --git a/cpp/src/gandiva/simple_arena.h b/cpp/src/gandiva/simple_arena.h index da00b33979d8..0c4d26961206 100644 --- a/cpp/src/gandiva/simple_arena.h +++ b/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,6 +92,9 @@ class SimpleArena { // List of allocated chunks. std::vector 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) @@ -95,11 +102,13 @@ inline SimpleArena::SimpleArena(arrow::MemoryPool* pool, int64_t min_chunk_size) 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;