diff --git a/cpp/src/gandiva/execution_context.h b/cpp/src/gandiva/execution_context.h index efa54687445..5e3d6f3fe3b 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/llvm_generator.cc b/cpp/src/gandiva/llvm_generator.cc index 7aa91c4dd6b..f787cf8e487 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,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 diff --git a/cpp/src/gandiva/simple_arena.h b/cpp/src/gandiva/simple_arena.h index da00b33979d..0c4d2696120 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; diff --git a/cpp/src/gandiva/tests/projector_test.cc b/cpp/src/gandiva/tests/projector_test.cc index 6e9450e1c65..afb0d503b18 100644 --- a/cpp/src/gandiva/tests/projector_test.cc +++ b/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(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"; +} + +// 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()); @@ -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(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 3d0ab51752e..34d39add6e4 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 d8181fe6751..beb92e626fd 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();