diff --git a/cpp/src/gandiva/decimal_ir.cc b/cpp/src/gandiva/decimal_ir.cc index 36e423e0647..5ece2bf9e81 100644 --- a/cpp/src/gandiva/decimal_ir.cc +++ b/cpp/src/gandiva/decimal_ir.cc @@ -243,12 +243,18 @@ DecimalIR::ValueWithOverflow DecimalIR::AddWithOverflowCheck(const ValueFull& x, return ValueWithOverflow(sum_descaled, overflow); } +std::pair DecimalIR::CreateEntryBlockReturnSlots() { + auto& entry = ir_builder()->GetInsertBlock()->getParent()->getEntryBlock(); + llvm::IRBuilder<> entry_builder(&entry, entry.begin()); + auto out_high_ptr = entry_builder.CreateAlloca(types()->i64_type(), nullptr, "out_hi"); + auto out_low_ptr = entry_builder.CreateAlloca(types()->i64_type(), nullptr, "out_low"); + return {out_high_ptr, out_low_ptr}; +} + // This is pretty complex, so use CPP fns. llvm::Value* DecimalIR::AddLarge(const ValueFull& x, const ValueFull& y, const ValueFull& out) { - auto block = ir_builder()->GetInsertBlock(); - auto out_high_ptr = new llvm::AllocaInst(types()->i64_type(), 0, "out_hi", block); - auto out_low_ptr = new llvm::AllocaInst(types()->i64_type(), 0, "out_low", block); + auto [out_high_ptr, out_low_ptr] = CreateEntryBlockReturnSlots(); auto x_split = ValueSplit::MakeFromInt128(this, x.value()); auto y_split = ValueSplit::MakeFromInt128(this, y.value()); @@ -454,9 +460,9 @@ llvm::Value* DecimalIR::CallDecimalFunction(const std::string& function_name, llvm::Value* result = nullptr; if (return_type == i128) { // for i128 ret, replace with two int64* args, and join them. - auto block = ir_builder()->GetInsertBlock(); - auto out_high_ptr = new llvm::AllocaInst(i64, 0, "out_hi", block); - auto out_low_ptr = new llvm::AllocaInst(i64, 0, "out_low", block); + // Decimal calls can be emitted inside the record loop or a CASE branch. + // Allocate return slots once per invocation, rather than once per row. + auto [out_high_ptr, out_low_ptr] = CreateEntryBlockReturnSlots(); dis_assembled_args.push_back(out_high_ptr); dis_assembled_args.push_back(out_low_ptr); diff --git a/cpp/src/gandiva/decimal_ir.h b/cpp/src/gandiva/decimal_ir.h index b11730f1e23..83f512b32f4 100644 --- a/cpp/src/gandiva/decimal_ir.h +++ b/cpp/src/gandiva/decimal_ir.h @@ -19,6 +19,7 @@ #include #include +#include #include #include "gandiva/function_ir_builder.h" @@ -137,6 +138,9 @@ class DecimalIR : public FunctionIRBuilder { ValueWithOverflow AddWithOverflowCheck(const ValueFull& x, const ValueFull& y, const ValueFull& out); + // Allocate decimal return slots in the current function's entry block. + std::pair CreateEntryBlockReturnSlots(); + // Do addition of large integers (both positive and negative). llvm::Value* AddLarge(const ValueFull& x, const ValueFull& y, const ValueFull& out); diff --git a/cpp/src/gandiva/tests/decimal_test.cc b/cpp/src/gandiva/tests/decimal_test.cc index 043bdc4605a..e395452d828 100644 --- a/cpp/src/gandiva/tests/decimal_test.cc +++ b/cpp/src/gandiva/tests/decimal_test.cc @@ -67,6 +67,42 @@ std::vector TestDecimal::MakeDecimalVector(std::vector return ret; } +TEST_F(TestDecimal, CastLargeBatch) { + // Allocating decimal return slots in the record loop can exhaust the stack + // when optimization is disabled. A single int32-to-decimal cast is sufficient. + constexpr int kNumRecords = 1 << 19; + auto input_field = arrow::field("large_batch_decimal", arrow::int32()); + auto schema = arrow::schema({input_field}); + auto decimal_type = arrow::decimal128(10, 0); + auto decimal = TreeExprBuilder::MakeFunction( + "castDECIMAL", {TreeExprBuilder::MakeField(input_field)}, decimal_type); + auto expression = + TreeExprBuilder::MakeExpression(decimal, arrow::field("result", decimal_type)); + + std::vector values(kNumRecords); + std::vector validity(kNumRecords); + std::vector expected_values(kNumRecords); + for (int i = 0; i < kNumRecords; ++i) { + values[i] = i % 17 - 8; + validity[i] = i % 11 != 0; + expected_values[i] = Decimal128(values[i]); + } + auto batch = arrow::RecordBatch::Make(schema, kNumRecords, + {MakeArrowArrayInt32(values, validity)}); + auto expected = MakeArrowArrayDecimal(decimal_type, expected_values, validity); + + for (bool optimize : {false, true}) { + SCOPED_TRACE(optimize); + auto config = std::make_shared(optimize); + std::shared_ptr projector; + ASSERT_OK(Projector::Make(schema, {expression}, config, &projector)); + arrow::ArrayVector output; + ASSERT_OK(projector->Evaluate(*batch, pool_, &output)); + ASSERT_EQ(output.size(), 1); + EXPECT_ARROW_ARRAY_EQUALS(expected, output[0]); + } +} + TEST_F(TestDecimal, TestSimple) { // schema for input fields constexpr int32_t precision = 36;