diff --git a/cpp/src/gandiva/context_helper.cc b/cpp/src/gandiva/context_helper.cc index 752fbea4c331..c60e952292d4 100644 --- a/cpp/src/gandiva/context_helper.cc +++ b/cpp/src/gandiva/context_helper.cc @@ -47,6 +47,13 @@ arrow::Status ExportedContextFunctions::AddMappings(Engine* engine) const { args, reinterpret_cast(gdv_fn_context_arena_malloc)); + // gdv_fn_context_arena_used_ptr + args = {types->i64_type()}; // int64_t context_ptr + + engine->AddGlobalMappingForFunc("gdv_fn_context_arena_used_ptr", types->i8_ptr_type(), + args, + reinterpret_cast(gdv_fn_context_arena_used_ptr)); + // gdv_fn_context_arena_reset args = {types->i64_type()}; // int64_t context_ptr @@ -70,6 +77,11 @@ uint8_t* gdv_fn_context_arena_malloc(int64_t context_ptr, int32_t size) { return context->arena()->Allocate(size); } +uint8_t* gdv_fn_context_arena_used_ptr(int64_t context_ptr) { + auto context = reinterpret_cast(context_ptr); + return context->arena()->used_ptr(); +} + void gdv_fn_context_arena_reset(int64_t context_ptr) { auto context = reinterpret_cast(context_ptr); return context->arena()->Reset(); diff --git a/cpp/src/gandiva/gdv_function_stubs.h b/cpp/src/gandiva/gdv_function_stubs.h index 4113f261ad76..28a71055694f 100644 --- a/cpp/src/gandiva/gdv_function_stubs.h +++ b/cpp/src/gandiva/gdv_function_stubs.h @@ -80,6 +80,8 @@ void gdv_fn_context_set_error_msg(int64_t context_ptr, const char* err_msg); GANDIVA_EXPORT uint8_t* gdv_fn_context_arena_malloc(int64_t context_ptr, int32_t data_len); +uint8_t* gdv_fn_context_arena_used_ptr(int64_t context_ptr); + void gdv_fn_context_arena_reset(int64_t context_ptr); bool in_expr_lookup_int32(int64_t ptr, int32_t value, bool in_validity); diff --git a/cpp/src/gandiva/llvm_generator.cc b/cpp/src/gandiva/llvm_generator.cc index 7aa91c4dd6b4..c2f0dbb183fe 100644 --- a/cpp/src/gandiva/llvm_generator.cc +++ b/cpp/src/gandiva/llvm_generator.cc @@ -394,6 +394,13 @@ Status LLVMGenerator::CodeGenExprValue(DexPtr value_expr, int buffer_count, // add jump to "loop block" at the end of the "setup block". builder->SetInsertPoint(loop_entry); + llvm::Value* arena_used_ptr = nullptr; + if (visitor.has_arena_allocs()) { + // Get the arena state pointer once per generated function instead of baking the + // ExecutionContext layout into the generated code. + arena_used_ptr = AddFunctionCall("gdv_fn_context_arena_used_ptr", + types()->i8_ptr_type(), {arg_context_ptr}); + } builder->CreateBr(loop_body); // save the value in the output vector. @@ -427,10 +434,23 @@ 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* arena_used = + builder->CreateLoad(types()->i8_type(), arena_used_ptr, "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 +1283,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 +1319,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 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; diff --git a/cpp/src/gandiva/tests/projector_test.cc b/cpp/src/gandiva/tests/projector_test.cc index 6e9450e1c650..dee184cad759 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)); } +// Before this fix, per-row arena scratch space was released after each row only if the +// expression contained a function that returned 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 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();