This is an automated email from the ASF dual-hosted git repository.
BiteTheDDDDt pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new c8b81647776 [improvement](be) Track Python UDF Arrow memory
allocations (#65361)
c8b81647776 is described below
commit c8b81647776bcd482b99665530f9fd5ba41f0d03
Author: Pxl <[email protected]>
AuthorDate: Wed Jul 15 23:25:28 2026 +0800
[improvement](be) Track Python UDF Arrow memory allocations (#65361)
Problem Summary: Python UDF/UDAF/UDTF input and request batch
construction used Arrow's default memory pool, so Arrow buffers
allocated while preparing Python function requests could bypass Doris
memory tracking.
This PR switches the synchronous Python UDF/UDAF/UDTF Arrow batch
construction paths to `ExecEnv::GetInstance()->arrow_memory_pool()`.
Benefits:
- Routes these Arrow buffer allocations through Doris' BE Arrow memory
pool, which is backed by Doris Allocator and memory tracker.
- Keeps Python function request/input batch memory accounting consistent
with other Doris Arrow integration paths.
- Avoids changing Python UDF/UDAF/UDTF semantics; this only changes the
allocator used for in-process Arrow buffers.
---
be/src/exprs/aggregate/aggregate_function_python_udaf.cpp | 11 +++++++----
be/src/exprs/function/function_python_udf.cpp | 4 +++-
be/src/exprs/table_function/python_udtf_function.cpp | 5 +++--
be/src/udf/python/python_udaf_client.cpp | 13 +++++++------
4 files changed, 20 insertions(+), 13 deletions(-)
diff --git a/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp
b/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp
index 4b6917f3962..cf27191ebe7 100644
--- a/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp
+++ b/be/src/exprs/aggregate/aggregate_function_python_udaf.cpp
@@ -33,6 +33,7 @@
#include "core/data_type/define_primitive_type.h"
#include "format/arrow/arrow_block_convertor.h"
#include "format/arrow/arrow_row_batch.h"
+#include "runtime/exec_env.h"
#include "runtime/user_function_cache.h"
#include "udf/python/python_env.h"
#include "udf/python/python_server.h"
@@ -66,8 +67,9 @@ Status AggregatePythonUDAFData::add(int64_t place_id, const
IColumn** columns,
std::shared_ptr<arrow::RecordBatch> batch;
// Zero-copy: convert only the specified range
- RETURN_IF_ERROR(convert_to_arrow_batch(input_block, schema,
arrow::default_memory_pool(),
- &batch, timezone_obj,
row_num_start, row_num_end));
+ RETURN_IF_ERROR(convert_to_arrow_batch(input_block, schema,
+
ExecEnv::GetInstance()->arrow_memory_pool(), &batch,
+ timezone_obj, row_num_start,
row_num_end));
// Send the batch (already sliced in convert_to_arrow_batch)
// Single place mode: no places column needed
RETURN_IF_ERROR(client->accumulate(place_id, true, *batch, 0,
batch->num_rows()));
@@ -111,8 +113,9 @@ Status AggregatePythonUDAFData::add_batch(AggregateDataPtr*
places, size_t place
std::shared_ptr<arrow::RecordBatch> batch;
// Zero-copy: convert only the [start, end) range
// This slice includes the places column automatically
- RETURN_IF_ERROR(convert_to_arrow_batch(input_block, schema,
arrow::default_memory_pool(),
- &batch, timezone_obj, start, end));
+ RETURN_IF_ERROR(convert_to_arrow_batch(input_block, schema,
+
ExecEnv::GetInstance()->arrow_memory_pool(), &batch,
+ timezone_obj, start, end));
// Send entire batch (already contains places column) to Python
// place_id=0 is ignored when is_single_place=false
RETURN_IF_ERROR(client->accumulate(0, false, *batch, 0, slice_rows));
diff --git a/be/src/exprs/function/function_python_udf.cpp
b/be/src/exprs/function/function_python_udf.cpp
index 54ea74ad456..ca9fa280f08 100644
--- a/be/src/exprs/function/function_python_udf.cpp
+++ b/be/src/exprs/function/function_python_udf.cpp
@@ -30,6 +30,7 @@
#include "common/status.h"
#include "core/block/block.h"
#include "format/arrow/arrow_block_convertor.h"
+#include "runtime/exec_env.h"
#include "runtime/user_function_cache.h"
#include "udf/python/python_server.h"
#include "udf/python/python_udf_client.h"
@@ -144,7 +145,8 @@ Status PythonFunctionCall::execute_impl(FunctionContext*
context, Block& block,
if (arguments.empty()) {
RETURN_IF_ERROR(make_zero_column_arrow_batch(schema, input_rows,
&input_batch));
} else {
- RETURN_IF_ERROR(convert_to_arrow_batch(input_block, schema,
arrow::default_memory_pool(),
+ RETURN_IF_ERROR(convert_to_arrow_batch(input_block, schema,
+
ExecEnv::GetInstance()->arrow_memory_pool(),
&input_batch, _timezone_obj));
}
RETURN_IF_ERROR(client->evaluate(*input_batch, &output_batch));
diff --git a/be/src/exprs/table_function/python_udtf_function.cpp
b/be/src/exprs/table_function/python_udtf_function.cpp
index eaaf4d5227a..fbbea4bfc81 100644
--- a/be/src/exprs/table_function/python_udtf_function.cpp
+++ b/be/src/exprs/table_function/python_udtf_function.cpp
@@ -38,6 +38,7 @@
#include "format/arrow/arrow_block_convertor.h"
#include "format/arrow/arrow_row_batch.h"
#include "format/arrow/arrow_utils.h"
+#include "runtime/exec_env.h"
#include "runtime/runtime_state.h"
#include "runtime/user_function_cache.h"
#include "udf/python/python_env.h"
@@ -141,8 +142,8 @@ Status PythonUDTFFunction::process_init(Block* block,
RuntimeState* state) {
RETURN_IF_ERROR(make_zero_column_arrow_batch(input_schema, input_rows,
&input_batch));
} else {
RETURN_IF_ERROR(convert_to_arrow_batch(input_block, input_schema,
- arrow::default_memory_pool(),
&input_batch,
- _timezone_obj));
+
ExecEnv::GetInstance()->arrow_memory_pool(),
+ &input_batch, _timezone_obj));
}
// Step 3: Call Python UDTF to evaluate all rows at once (similar to Java
UDTF's JNI call)
diff --git a/be/src/udf/python/python_udaf_client.cpp
b/be/src/udf/python/python_udaf_client.cpp
index 4308750ccec..d334017415b 100644
--- a/be/src/udf/python/python_udaf_client.cpp
+++ b/be/src/udf/python/python_udaf_client.cpp
@@ -30,6 +30,7 @@
#include "common/compiler_util.h"
#include "common/status.h"
#include "format/arrow/arrow_utils.h"
+#include "runtime/exec_env.h"
#include "udf/python/python_udf_meta.h"
#include "udf/python/python_udf_runtime.h"
#include "util/unaligned.h"
@@ -509,7 +510,7 @@ Status PythonUDAFClient::_create_data_request_batch(const
arrow::RecordBatch& in
columns.push_back(input_data.column(num_input_columns - 1));
} else {
// Create NULL places column for single-place mode
- arrow::Int64Builder places_builder;
+ arrow::Int64Builder
places_builder(ExecEnv::GetInstance()->arrow_memory_pool());
std::shared_ptr<arrow::Array> places_array;
RETURN_DORIS_STATUS_IF_ERROR(places_builder.AppendNulls(input_data.num_rows()));
RETURN_DORIS_STATUS_IF_ERROR(places_builder.Finish(&places_array));
@@ -517,7 +518,7 @@ Status PythonUDAFClient::_create_data_request_batch(const
arrow::RecordBatch& in
}
// Add NULL binary_data column
- arrow::BinaryBuilder binary_builder;
+ arrow::BinaryBuilder
binary_builder(ExecEnv::GetInstance()->arrow_memory_pool());
std::shared_ptr<arrow::Array> binary_array;
RETURN_DORIS_STATUS_IF_ERROR(binary_builder.AppendNulls(input_data.num_rows()));
RETURN_DORIS_STATUS_IF_ERROR(binary_builder.Finish(&binary_array));
@@ -538,7 +539,7 @@ Status PythonUDAFClient::_create_binary_request_batch(
for (int i = 0; i < num_data_columns; ++i) {
std::unique_ptr<arrow::ArrayBuilder> builder;
std::shared_ptr<arrow::Array> null_array;
-
RETURN_DORIS_STATUS_IF_ERROR(arrow::MakeBuilder(arrow::default_memory_pool(),
+
RETURN_DORIS_STATUS_IF_ERROR(arrow::MakeBuilder(ExecEnv::GetInstance()->arrow_memory_pool(),
_schema->field(i)->type(), &builder));
RETURN_DORIS_STATUS_IF_ERROR(builder->AppendNull());
RETURN_DORIS_STATUS_IF_ERROR(builder->Finish(&null_array));
@@ -546,7 +547,7 @@ Status PythonUDAFClient::_create_binary_request_batch(
}
// Create binary_data column
- arrow::BinaryBuilder binary_builder;
+ arrow::BinaryBuilder
binary_builder(ExecEnv::GetInstance()->arrow_memory_pool());
std::shared_ptr<arrow::Array> binary_array;
RETURN_DORIS_STATUS_IF_ERROR(
binary_builder.Append(binary_data->data(),
static_cast<int32_t>(binary_data->size())));
@@ -571,8 +572,8 @@ Status
PythonUDAFClient::_get_empty_request_batch(std::shared_ptr<arrow::RecordB
auto field = _schema->field(i);
std::unique_ptr<arrow::ArrayBuilder> builder;
std::shared_ptr<arrow::Array> null_array;
- RETURN_DORIS_STATUS_IF_ERROR(
- arrow::MakeBuilder(arrow::default_memory_pool(),
field->type(), &builder));
+
RETURN_DORIS_STATUS_IF_ERROR(arrow::MakeBuilder(ExecEnv::GetInstance()->arrow_memory_pool(),
+ field->type(),
&builder));
RETURN_DORIS_STATUS_IF_ERROR(builder->AppendNull());
RETURN_DORIS_STATUS_IF_ERROR(builder->Finish(&null_array));
columns.push_back(null_array);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
