This is an automated email from the ASF dual-hosted git repository.
HappenLee 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 e870abeee7c [Fix](test) make pyudf ut deterministic and meaningful
(#67352)
e870abeee7c is described below
commit e870abeee7c4ec066143261882887015149e7f69
Author: linrrarity <[email protected]>
AuthorDate: Thu Sep 10 18:23:59 2026 +0800
[Fix](test) make pyudf ut deterministic and meaningful (#67352)
Problem Summary:
Some Python UDF BE unit tests were flaky and provided limited
verification value. Several tests launched shell-based fake Python
processes that only created Unix socket files or stayed alive for a
fixed duration. These tests exercised process-management branches but
did not verify that Doris could start the production Python server,
establish an Arrow Flight connection, or execute an actual UDF.
This PR:
- Removes or rewrites low-value Python UDF tests that only exercised
implementation details or line coverage.
- Keeps meaningful coverage for process startup failures, shutdown,
recovery, pool initialization, and concurrent process selection.
- Fixes the race in the least-shared-process test by waiting until the
complete process pool is available before checking process distribution.
- Adds an end-to-end BE unit test that:
- Creates a Python runtime through the production venv scanner.
- Starts the real `python_server.py`.
- Connects through Arrow Flight.
- Executes a nullable inline scalar UDF.
- Verifies the returned values and null propagation.
- Adds `PYTHON_UDF_TEST_INTERPRETER` as the only Python interpreter
configuration for this test. It defaults to `/usr/bin/python`, matching
the community P0 pipeline, and can be overridden locally in
`custom_env.sh`.
### Release note
None
### Check List (For Author)
- Test <!-- At least one of them must be included. -->
- [ ] Regression test
- [ ] Unit Test
- [ ] Manual test (add detailed scripts or steps below)
- [ ] No need to test or manual test. Explain why:
- [ ] This is a refactor/code format and no logic has been changed.
- [ ] Previous test can cover this change.
- [ ] No code files have been changed.
- [ ] Other reason <!-- Add your reason? -->
- Behavior changed:
- [ ] No.
- [ ] Yes. <!-- Explain the behavior change -->
- Does this need documentation?
- [ ] No.
- [ ] Yes. <!-- Add document PR link here. eg:
https://github.com/apache/doris-website/pull/1214 -->
### Check List (For Reviewer who merge this PR)
- [ ] Confirm the release note
- [ ] Confirm test cases
- [ ] Confirm document
- [ ] Add branch pick label <!-- Add branch pick label that this PR
should merge into -->
---
be/src/udf/python/python_server.cpp | 55 +++-
be/src/udf/python/python_server.h | 8 +-
be/test/udf/python/python_env_test.cpp | 179 ++++-------
be/test/udf/python/python_server_test.cpp | 398 +++++++++++++------------
be/test/udf/python/python_udf_meta_test.cpp | 144 ++++-----
be/test/udf/python/python_udf_runtime_test.cpp | 225 ++------------
6 files changed, 397 insertions(+), 612 deletions(-)
diff --git a/be/src/udf/python/python_server.cpp
b/be/src/udf/python/python_server.cpp
index cf78591a2ef..2d06caa3110 100644
--- a/be/src/udf/python/python_server.cpp
+++ b/be/src/udf/python/python_server.cpp
@@ -36,6 +36,7 @@
#include "arrow/flight/client.h"
#include "common/config.h"
#include "common/status.h"
+#include "cpp/sync_point.h"
#include "runtime/thread_context.h"
#include "udf/python/python_udaf_client.h"
#include "udf/python/python_udf_client.h"
@@ -91,17 +92,37 @@ std::vector<ProcessPtr>
PythonServerManager::process_pool_snapshot_for_test(
return versioned_pool->processes;
}
-bool PythonServerManager::process_pool_is_initializing_for_test(const
PythonVersion& version) {
- auto versioned_pool = _get_or_create_process_pool(version).value();
- std::lock_guard<std::mutex> lock(versioned_pool->mutex);
- return versioned_pool->state == PoolState::INITIALIZING;
-}
-
bool PythonServerManager::process_pool_is_initialized_for_test(const
PythonVersion& version) {
auto versioned_pool = _get_or_create_process_pool(version).value();
std::lock_guard<std::mutex> lock(versioned_pool->mutex);
return versioned_pool->state == PoolState::INITIALIZED;
}
+
+bool PythonServerManager::wait_for_process_pool_initialized_for_test(
+ const PythonVersion& version, std::chrono::milliseconds timeout) {
+ auto versioned_pool_result = _get_or_create_process_pool(version);
+ if (!versioned_pool_result.has_value()) {
+ return false;
+ }
+ auto versioned_pool = versioned_pool_result.value();
+ std::unique_lock<std::mutex> lock(versioned_pool->mutex);
+ return versioned_pool->cv.wait_for(lock, timeout, [&versioned_pool]() {
+ return versioned_pool->state == PoolState::INITIALIZED;
+ });
+}
+
+bool
PythonServerManager::wait_for_process_pool_initialization_finished_for_test(
+ const PythonVersion& version, std::chrono::milliseconds timeout) {
+ auto versioned_pool_result = _get_or_create_process_pool(version);
+ if (!versioned_pool_result.has_value()) {
+ return false;
+ }
+ auto versioned_pool = versioned_pool_result.value();
+ std::unique_lock<std::mutex> lock(versioned_pool->mutex);
+ return versioned_pool->cv.wait_for(lock, timeout, [&versioned_pool]() {
+ return versioned_pool->state != PoolState::INITIALIZING;
+ });
+}
#endif
bool PythonServerManager::_select_alive_process_from_pool(const
std::vector<ProcessPtr>& pool,
@@ -150,6 +171,10 @@ PythonServerManager::_ensure_pool_initialized(const
PythonVersion& version) {
auto versioned_pool = versioned_pool_result.value();
const int max_pool_size = config::max_python_process_num > 0 ?
config::max_python_process_num
:
CpuInfo::num_cores();
+ auto process_pool_init_timeout = PROCESS_POOL_INIT_TIMEOUT;
+ TEST_SYNC_POINT_CALLBACK(
+
"PythonServerManager::_ensure_pool_initialized:process_pool_init_timeout",
+ &process_pool_init_timeout);
std::unique_lock<std::mutex> lock(versioned_pool->mutex);
if (versioned_pool->state == PoolState::INITIALIZED) {
@@ -162,17 +187,21 @@ PythonServerManager::_ensure_pool_initialized(const
PythonVersion& version) {
versioned_pool->has_available_process = false;
versioned_pool->processes.resize(max_pool_size);
auto init_finished_count = std::make_shared<std::atomic<int>>(0);
+ TEST_SYNC_POINT_CALLBACK(
+
"PythonServerManager::_ensure_pool_initialized:generation_started", &version,
+ init_finished_count.get());
LOG(INFO) << "Initializing Python process pool for version " <<
version.to_string()
<< " with " << max_pool_size << " processes
(config::max_python_process_num="
<< config::max_python_process_num << ", CPU cores=" <<
CpuInfo::num_cores()
<< ")";
- std::thread([this, versioned_pool, init_finished_count,
max_pool_size]() {
+ std::thread([this, versioned_pool, init_finished_count,
max_pool_size,
+ process_pool_init_timeout]() {
SCOPED_INIT_THREAD_CONTEXT();
std::unique_lock<std::mutex> lock(versioned_pool->mutex);
versioned_pool->cv.wait_for(
- lock, PROCESS_POOL_INIT_TIMEOUT,
+ lock, process_pool_init_timeout,
[&versioned_pool, init_finished_count,
max_pool_size]() {
return versioned_pool->state !=
PoolState::INITIALIZING ||
init_finished_count->load(std::memory_order_acquire) >=
@@ -223,13 +252,16 @@ PythonServerManager::_ensure_pool_initialized(const
PythonVersion& version) {
if (process_to_shutdown) {
process_to_shutdown->shutdown();
}
+ TEST_SYNC_POINT_CALLBACK(
+
"PythonServerManager::_ensure_pool_initialized:init_worker_finished",
+ &version, init_finished_count.get());
}).detach();
}
}
// Wait only for the first usable process. INITIALIZED is set later by
the last init worker
// after every slot has attempted initialization.
- versioned_pool->cv.wait_for(lock, PROCESS_POOL_INIT_TIMEOUT,
[&versioned_pool]() {
+ versioned_pool->cv.wait_for(lock, process_pool_init_timeout,
[&versioned_pool]() {
return versioned_pool->has_available_process ||
versioned_pool->state == PoolState::STOPPED ||
versioned_pool->state != PoolState::INITIALIZING;
@@ -243,7 +275,7 @@ PythonServerManager::_ensure_pool_initialized(const
PythonVersion& version) {
return ResultError(Status::Error<ErrorCode::SERVICE_UNAVAILABLE>(
"Failed to initialize Python process pool for version {}: no
process became available "
"within {} ms",
- version.to_string(), PROCESS_POOL_INIT_TIMEOUT.count()));
+ version.to_string(), process_pool_init_timeout.count()));
}
Status PythonServerManager::_get_process(
@@ -346,6 +378,9 @@ Status PythonServerManager::fork(const PythonVersion&
version, ProcessPtr* proce
boost::process::environment env = boost::this_process::environment();
boost::process::ipstream child_output;
+ // Test synchronization belongs before child creation so it cannot consume
readiness timeout.
+ TEST_SYNC_POINT_CALLBACK("PythonServerManager::fork:before_process_start",
&version);
+
try {
boost::process::child c(
python_executable_path, args, boost::process::std_out >
child_output,
diff --git a/be/src/udf/python/python_server.h
b/be/src/udf/python/python_server.h
index 5ee763d5f52..c72a2db43d1 100644
--- a/be/src/udf/python/python_server.h
+++ b/be/src/udf/python/python_server.h
@@ -67,10 +67,14 @@ public:
std::vector<ProcessPtr> process_pool_snapshot_for_test(const
PythonVersion& version);
- bool process_pool_is_initializing_for_test(const PythonVersion& version);
-
bool process_pool_is_initialized_for_test(const PythonVersion& version);
+ bool wait_for_process_pool_initialized_for_test(const PythonVersion&
version,
+ std::chrono::milliseconds
timeout);
+
+ bool wait_for_process_pool_initialization_finished_for_test(const
PythonVersion& version,
+
std::chrono::milliseconds timeout);
+
Status broadcast_action_to_processes_for_test(const std::string&
action_type,
const std::string& body,
const std::string& log_name)
{
diff --git a/be/test/udf/python/python_env_test.cpp
b/be/test/udf/python/python_env_test.cpp
index 361d3c59abc..4e48ee71d7a 100644
--- a/be/test/udf/python/python_env_test.cpp
+++ b/be/test/udf/python/python_env_test.cpp
@@ -23,6 +23,7 @@
#include <filesystem>
#include <fstream>
#include <string>
+#include <unordered_map>
#include <vector>
namespace doris {
@@ -76,37 +77,6 @@ protected:
// PythonVersion tests
// ============================================================================
-TEST_F(PythonEnvTest, PythonVersionDefaultConstruction) {
- PythonVersion pv;
- EXPECT_TRUE(pv.full_version.empty());
- EXPECT_TRUE(pv.base_path.empty());
- EXPECT_TRUE(pv.executable_path.empty());
-}
-
-TEST_F(PythonEnvTest, PythonVersionExplicitConstruction) {
- PythonVersion pv("3.9.16", "/opt/python", "/opt/python/bin/python3");
- EXPECT_EQ(pv.full_version, "3.9.16");
- EXPECT_EQ(pv.base_path, "/opt/python");
- EXPECT_EQ(pv.executable_path, "/opt/python/bin/python3");
-}
-
-TEST_F(PythonEnvTest, PythonVersionEquality) {
- PythonVersion pv1("3.9.16", "/opt/python", "/opt/python/bin/python3");
- PythonVersion pv2("3.9.16", "/opt/python", "/opt/python/bin/python3");
- PythonVersion pv3("3.10.0", "/opt/python", "/opt/python/bin/python3");
- PythonVersion pv4("3.9.16", "/different/path", "/opt/python/bin/python3");
-
- EXPECT_EQ(pv1, pv2);
- EXPECT_FALSE(pv1 == pv3);
- EXPECT_FALSE(pv1 == pv4);
-}
-
-TEST_F(PythonEnvTest, PythonVersionGetters) {
- PythonVersion pv("3.9.16", "/opt/python", "/opt/python/bin/python3");
- EXPECT_EQ(pv.get_base_path(), "/opt/python");
- EXPECT_EQ(pv.get_executable_path(), "/opt/python/bin/python3");
-}
-
TEST_F(PythonEnvTest, PythonVersionIsValidEmpty) {
PythonVersion pv;
EXPECT_FALSE(pv.is_valid());
@@ -138,50 +108,34 @@ TEST_F(PythonEnvTest,
PythonVersionIsValidWithExistingPaths) {
fs::permissions(exec_path, fs::perms::owner_all);
PythonVersion pv("3.9.16", base_path, exec_path);
- // is_valid() also checks that extract_python_version works, which
requires a real python
- // So this will fail on fake executable, but we verify paths exist
- EXPECT_TRUE(fs::exists(base_path));
- EXPECT_TRUE(fs::exists(exec_path));
-}
-
-TEST_F(PythonEnvTest, PythonVersionToString) {
- PythonVersion pv("3.9.16", "/opt/python", "/opt/python/bin/python3");
- std::string str = pv.to_string();
- EXPECT_TRUE(str.find("3.9.16") != std::string::npos);
- EXPECT_TRUE(str.find("/opt/python") != std::string::npos);
- EXPECT_TRUE(str.find("/opt/python/bin/python3") != std::string::npos);
-}
-
-TEST_F(PythonEnvTest, PythonVersionHash) {
- PythonVersion pv1("3.9.16", "/opt/python", "/opt/python/bin/python3");
- PythonVersion pv2("3.9.16", "/opt/python", "/opt/python/bin/python3");
- PythonVersion pv3("3.10.0", "/opt/python", "/opt/python/bin/python3");
-
- std::hash<PythonVersion> hasher;
- EXPECT_EQ(hasher(pv1), hasher(pv2));
- // Different versions should (very likely) have different hashes
- EXPECT_NE(hasher(pv1), hasher(pv3));
+ EXPECT_TRUE(pv.is_valid());
}
-// ============================================================================
-// PythonEnvironment tests
-// ============================================================================
-
-TEST_F(PythonEnvTest, PythonEnvironmentConstruction) {
- PythonVersion pv("3.9.16", "/opt/python", "/opt/python/bin/python3");
- PythonEnvironment env("test_env", pv);
+TEST_F(PythonEnvTest, PythonVersionWorksAsProcessPoolMapKey) {
+ PythonVersion baseline("3.9.16", "/opt/python", "/opt/python/bin/python3");
+ PythonVersion equal_key("3.9.16", "/opt/python",
"/opt/python/bin/python3");
+ PythonVersion different_version("3.10.0", "/opt/python",
"/opt/python/bin/python3");
+ PythonVersion different_base("3.9.16", "/different/path",
"/opt/python/bin/python3");
+ PythonVersion different_executable("3.9.16", "/opt/python",
"/different/path/bin/python3");
- EXPECT_EQ(env.env_name, "test_env");
- EXPECT_EQ(env.python_version.full_version, "3.9.16");
-}
+ // PythonServerManager keys process pools by the complete PythonVersion.
Equal keys must find
+ // the same pool, while a change in any identity field must select a
different pool.
+ EXPECT_EQ(baseline, equal_key);
+ EXPECT_NE(baseline, different_version);
+ EXPECT_NE(baseline, different_base);
+ EXPECT_NE(baseline, different_executable);
-TEST_F(PythonEnvTest, PythonEnvironmentToString) {
- PythonVersion pv("3.9.16", "/opt/python", "/opt/python/bin/python3");
- PythonEnvironment env("test_env", pv);
+ std::unordered_map<PythonVersion, int> pools;
+ pools.emplace(baseline, 1);
+ pools.emplace(different_version, 2);
+ pools.emplace(different_base, 3);
+ pools.emplace(different_executable, 4);
- std::string str = env.to_string();
- EXPECT_TRUE(str.find("test_env") != std::string::npos);
- EXPECT_TRUE(str.find("3.9.16") != std::string::npos);
+ ASSERT_EQ(pools.size(), 4);
+ EXPECT_EQ(pools.at(equal_key), 1);
+ EXPECT_EQ(pools.at(different_version), 2);
+ EXPECT_EQ(pools.at(different_base), 3);
+ EXPECT_EQ(pools.at(different_executable), 4);
}
TEST_F(PythonEnvTest, PythonEnvironmentIsValidWithInvalidVersion) {
@@ -231,6 +185,44 @@ TEST_F(PythonEnvTest,
ScanFromVenvRootPathNonExistentInterpreter) {
EXPECT_TRUE(status.to_string().find("Interpreter path not found") !=
std::string::npos);
}
+TEST_F(PythonEnvTest, ScanFromVenvRootPathCreatesAndValidatesEnvironment) {
+ const std::string root_path = test_dir_ + "/venv_root";
+ const std::string interpreter_path = test_dir_ + "/python3.11";
+ fs::create_directories(root_path);
+
+ // Model the two interpreter operations used by the production VENV path:
version discovery
+ // and `python -m venv`. The generated environment also answers --version,
so the scanner must
+ // complete its normal validation instead of accepting a directory shape
alone.
+ {
+ std::ofstream ofs(interpreter_path);
+ ofs << "#!/bin/bash\n";
+ ofs << "if [ \"$1\" = \"--version\" ]; then\n";
+ ofs << " echo 'Python 3.11.9'\n";
+ ofs << " exit 0\n";
+ ofs << "fi\n";
+ ofs << "if [ \"$1\" = \"-m\" ] && [ \"$2\" = \"venv\" ] && "
+ "[ \"$3\" = \"--system-site-packages\" ]; then\n";
+ ofs << " mkdir -p \"$4/bin\" \"$4/lib/python3.11/site-packages\"\n";
+ ofs << " cp \"$0\" \"$4/bin/python\"\n";
+ ofs << " exit 0\n";
+ ofs << "fi\n";
+ ofs << "exit 1\n";
+ }
+ fs::permissions(interpreter_path, fs::perms::owner_all);
+
+ std::vector<PythonEnvironment> envs;
+ Status status =
+ PythonEnvironment::scan_from_venv_root_path(root_path,
{interpreter_path}, &envs);
+
+ ASSERT_TRUE(status.ok()) << status.to_string();
+ ASSERT_EQ(envs.size(), 1);
+ EXPECT_EQ(envs[0].env_name, "python3.11.9");
+ EXPECT_EQ(envs[0].python_version.full_version, "3.11.9");
+ EXPECT_EQ(envs[0].python_version.base_path, root_path + "/python3.11.9");
+ EXPECT_EQ(envs[0].python_version.executable_path, root_path +
"/python3.11.9/bin/python");
+ EXPECT_TRUE(fs::is_directory(root_path +
"/python3.11.9/lib/python3.11/site-packages"));
+}
+
// ============================================================================
// PythonEnvScanner tests
// ============================================================================
@@ -251,38 +243,12 @@ TEST_F(PythonEnvTest, PythonEnvScannerGetVersionNotFound)
{
EXPECT_TRUE(status.to_string().find("not found available version") !=
std::string::npos);
}
-TEST_F(PythonEnvTest, CondaEnvScannerProperties) {
- CondaEnvScanner scanner(test_dir_);
- EXPECT_EQ(scanner.env_type(), PythonEnvType::CONDA);
- EXPECT_EQ(scanner.root_path(), test_dir_);
-}
-
-TEST_F(PythonEnvTest, CondaEnvScannerToString) {
- CondaEnvScanner scanner(test_dir_);
- std::string str = scanner.to_string();
- EXPECT_TRUE(str.find("Conda environments") != std::string::npos);
-}
-
TEST_F(PythonEnvTest, CondaEnvScannerScanNonExistent) {
CondaEnvScanner scanner("/non/existent/path");
Status status = scanner.scan();
EXPECT_FALSE(status.ok());
}
-TEST_F(PythonEnvTest, VenvEnvScannerProperties) {
- std::vector<std::string> paths = {"/usr/bin/python3"};
- VenvEnvScanner scanner(test_dir_, paths);
- EXPECT_EQ(scanner.env_type(), PythonEnvType::VENV);
- EXPECT_EQ(scanner.root_path(), test_dir_);
-}
-
-TEST_F(PythonEnvTest, VenvEnvScannerToString) {
- std::vector<std::string> paths = {"/usr/bin/python3"};
- VenvEnvScanner scanner(test_dir_, paths);
- std::string str = scanner.to_string();
- EXPECT_TRUE(str.find("Venv environments") != std::string::npos);
-}
-
TEST_F(PythonEnvTest, VenvEnvScannerScanNonExistentInterpreter) {
std::vector<std::string> paths = {"/non/existent/python3"};
VenvEnvScanner scanner(test_dir_, paths);
@@ -327,25 +293,6 @@ TEST_F(PythonEnvTest,
PythonVersionManagerInitUnsupportedType) {
// Tests covering additional paths using fake Python scripts
// ============================================================================
-TEST_F(PythonEnvTest, PythonVersionIsValidWithFakePython) {
- // Create a fake Python script that prints "Python 3.9.16"
- std::string base_path = test_dir_ + "/fake_python";
- std::string bin_path = base_path + "/bin";
- std::string exec_path = bin_path + "/python3";
- fs::create_directories(bin_path);
-
- {
- std::ofstream ofs(exec_path);
- ofs << "#!/bin/bash\n";
- ofs << "echo 'Python 3.9.16'\n";
- }
- fs::permissions(exec_path, fs::perms::owner_all);
-
- PythonVersion pv("3.9.16", base_path, exec_path);
- // Now is_valid() should pass because the version matches
- EXPECT_TRUE(pv.is_valid());
-}
-
TEST_F(PythonEnvTest, PythonVersionIsValidVersionMismatch) {
// Create a fake Python script that prints a different version
std::string base_path = test_dir_ + "/fake_python2";
@@ -483,7 +430,7 @@ TEST_F(PythonEnvTest, ScanFromCondaRootPathMultipleEnvs) {
EXPECT_EQ(envs.size(), 2);
}
-TEST_F(PythonEnvTest, ScanFromCondaRootPathSkipsInvalidEnvs) {
+TEST_F(PythonEnvTest, ScanFromCondaRootPathFailsOnInvalidEnv) {
// Create a valid and an invalid environment
std::string conda_root = test_dir_ + "/conda_mixed";
diff --git a/be/test/udf/python/python_server_test.cpp
b/be/test/udf/python/python_server_test.cpp
index 3ec463bbe1c..39e6143bc14 100644
--- a/be/test/udf/python/python_server_test.cpp
+++ b/be/test/udf/python/python_server_test.cpp
@@ -26,14 +26,17 @@
#include <filesystem>
#include <fstream>
#include <future>
+#include <optional>
#include <string>
#include <vector>
#include "common/config.h"
#include "common/status.h"
+#include "cpp/sync_point.h"
#include "udf/python/python_env.h"
#include "udf/python/python_udf_client.h"
#include "udf/python/python_udf_meta.h"
+#include "util/defer_op.h"
namespace doris {
@@ -43,7 +46,7 @@ namespace bp = boost::process;
class PythonServerTest : public ::testing::Test {
protected:
std::string test_dir_;
- const char* original_doris_home_ = nullptr;
+ std::optional<std::string> original_doris_home_;
int original_max_python_process_num_ = 0;
void SetUp() override {
@@ -51,7 +54,9 @@ protected:
std::to_string(getpid()) + "_" + std::to_string(rand());
fs::create_directories(test_dir_);
- original_doris_home_ = std::getenv("DORIS_HOME");
+ if (const char* doris_home = std::getenv("DORIS_HOME")) {
+ original_doris_home_ = doris_home;
+ }
original_max_python_process_num_ = config::max_python_process_num;
}
@@ -64,7 +69,7 @@ protected:
}
if (original_doris_home_) {
- setenv("DORIS_HOME", original_doris_home_, 1);
+ setenv("DORIS_HOME", original_doris_home_->c_str(), 1);
} else {
unsetenv("DORIS_HOME");
}
@@ -86,6 +91,7 @@ protected:
ofs << " echo 'Python " << version << "'\n";
ofs << " exit 0\n";
ofs << "fi\n";
+ ofs << "if [ ! -f \"$2\" ]; then exit 2; fi\n";
// Extract socket path prefix from args and create the socket file
// Arg format: -u script.py grpc+unix:///tmp/doris_python_udf
ofs << "SOCKET_PREFIX=\"$3\"\n";
@@ -103,32 +109,6 @@ protected:
return python_path;
}
- std::string create_fake_python_with_delay_and_socket_creation(const
std::string& binary_name,
- const
std::string& version,
- int
delay_ms) {
- std::string bin_dir = test_dir_ + "/bin";
- std::string python_path = bin_dir + "/" + binary_name;
- fs::create_directories(bin_dir);
-
- std::ofstream ofs(python_path);
- ofs << "#!/bin/bash\n";
- ofs << "if [ \"$1\" = \"--version\" ]; then\n";
- ofs << " echo 'Python " << version << "'\n";
- ofs << " exit 0\n";
- ofs << "fi\n";
- ofs << "sleep " << (delay_ms / 1000.0) << "\n";
- ofs << "SOCKET_PREFIX=\"$3\"\n";
- ofs << "SOCKET_BASE=\"${SOCKET_PREFIX#grpc+unix://}\"\n";
- ofs << "SOCKET_FILE=\"${SOCKET_BASE}_$$.sock\"\n";
- ofs << "touch \"$SOCKET_FILE\"\n";
- ofs << "trap 'rm -f \"$SOCKET_FILE\"; exit 0' TERM INT\n";
- ofs << "while true; do sleep 1; done\n";
- ofs.close();
- fs::permissions(python_path, fs::perms::owner_all);
-
- return python_path;
- }
-
std::string create_fake_python_without_socket_creation(const std::string&
binary_name,
const std::string&
version) {
std::string bin_dir = test_dir_ + "/bin";
@@ -141,6 +121,7 @@ protected:
ofs << " echo 'Python " << version << "'\n";
ofs << " exit 0\n";
ofs << "fi\n";
+ ofs << "if [ ! -f \"$2\" ]; then exit 2; fi\n";
ofs << "trap '' TERM\n";
ofs << "while true; do sleep 1; done\n";
ofs.close();
@@ -162,6 +143,7 @@ protected:
ofs << " echo 'Python " << version << "'\n";
ofs << " exit 0\n";
ofs << "fi\n";
+ ofs << "if [ ! -f \"$2\" ]; then exit 2; fi\n";
ofs << "if mkdir \"" << first_start_dir << "\" 2>/dev/null; then\n";
ofs << " trap '' TERM\n";
ofs << " while true; do sleep 1; done\n";
@@ -226,32 +208,6 @@ TEST_F(PythonServerTest, SingletonReturnsSameInstance) {
EXPECT_EQ(&mgr1, &mgr2);
}
-// ============================================================================
-// PythonServerManager::_get_process() - process retrieval test
-// ============================================================================
-
-TEST_F(PythonServerTest, EnsurePoolInitializedCanInitializeEmptyPoolForTest) {
- PythonServerManager mgr;
-
- setup_doris_home();
- std::string python_path =
create_fake_python_with_socket_creation("3.9.16");
- PythonVersion version("3.9.16", test_dir_, python_path);
- config::max_python_process_num = 1;
-
- mgr.set_process_pool_for_test(version, {}, false);
- auto pool_result = mgr._ensure_pool_initialized(version);
- ASSERT_TRUE(pool_result.has_value()) << pool_result.error().to_string();
-
- ProcessPtr process;
- Status status = mgr._get_process(version, pool_result.value(), &process);
-
- EXPECT_TRUE(status.ok()) << status.to_string();
- ASSERT_NE(process, nullptr);
- EXPECT_TRUE(process->is_alive());
-
- mgr.shutdown();
-}
-
// ============================================================================
// PythonServerManager::fork() - process creation test
// ============================================================================
@@ -271,25 +227,21 @@ TEST_F(PythonServerTest,
ForkWithNonExistentPythonReturnsError) {
TEST_F(PythonServerTest, ForkWithMissingFlightServerReturnsError) {
PythonServerManager mgr;
-
- // Set DORIS_HOME to test directory (no flight server script)
- setenv("DORIS_HOME", test_dir_.c_str(), 1);
-
- // Create a fake python executable
- std::string python_path = test_dir_ + "/bin/python3";
- fs::create_directories(test_dir_ + "/bin");
- {
- std::ofstream ofs(python_path);
- ofs << "#!/bin/bash\nexit 1"; // exits immediately
- }
- fs::permissions(python_path, fs::perms::owner_all);
-
+ setup_doris_home();
+ std::string python_path =
create_fake_python_with_socket_creation("3.9.16");
PythonVersion version("3.9.16", test_dir_, python_path);
+ // Prove that the same executable and arguments start correctly while the
production server
+ // entry exists, then remove only that entry to isolate the failure cause.
+ ProcessPtr healthy_process;
+ ASSERT_TRUE(mgr.fork(version, &healthy_process).ok());
+ ASSERT_NE(healthy_process, nullptr);
+ healthy_process->shutdown();
+
+ ASSERT_TRUE(fs::remove(test_dir_ +
"/plugins/python_udf/python_server.py"));
ProcessPtr process;
Status status = mgr.fork(version, &process);
- // Verify: when the flight server script does not exist, fork should fail
EXPECT_FALSE(status.ok());
EXPECT_EQ(process, nullptr);
}
@@ -401,19 +353,86 @@ TEST_F(PythonServerTest,
EnsurePoolInitializedWithInvalidVersionFails) {
result.error().to_string().find("Timed out") !=
std::string::npos);
}
-TEST_F(PythonServerTest,
EnsurePoolInitializedReturnsImmediatelyWhenAllWorkersFail) {
+TEST_F(PythonServerTest, EnsurePoolInitializedRetriesAfterRuntimeIsRepaired) {
+ setup_doris_home();
+ config::max_python_process_num = 1;
+
+ // The first executable starts but never publishes its Flight socket. This
reproduces a real
+ // worker-start failure and verifies that the pool does not remain stuck
in INITIALIZING.
+ std::string python_path =
create_fake_python_without_socket_creation("python3", "3.9.16");
+ PythonVersion version("3.9.16", test_dir_, python_path);
PythonServerManager mgr;
- config::max_python_process_num = 2;
- PythonVersion invalid_version("3.9.16", test_dir_, test_dir_ +
"/missing_python");
+ struct GenerationCompletionLatch {
+ std::mutex mutex;
+ std::condition_variable cv;
+ std::atomic<int>* generation = nullptr;
+ bool worker_finished = false;
+ };
+ auto generation_completion = std::make_shared<GenerationCompletionLatch>();
+ auto* sync_point = SyncPoint::get_instance();
+ Defer clear_sync_point {[sync_point]() {
+ sync_point->disable_processing();
+ sync_point->clear_call_back(
+
"PythonServerManager::_ensure_pool_initialized:generation_started");
+ sync_point->clear_call_back(
+
"PythonServerManager::_ensure_pool_initialized:init_worker_finished");
+ }};
+ sync_point->set_call_back(
+ "PythonServerManager::_ensure_pool_initialized:generation_started",
+ [generation_completion, version](auto&& args) {
+ const auto* started_version = try_any_cast<const
PythonVersion*>(args.at(0));
+ if (*started_version != version) {
+ return;
+ }
+ std::lock_guard lock(generation_completion->mutex);
+ generation_completion->generation =
try_any_cast<std::atomic<int>*>(args.at(1));
+ });
+ sync_point->set_call_back(
+
"PythonServerManager::_ensure_pool_initialized:init_worker_finished",
+ [generation_completion, version](auto&& args) {
+ const auto* finished_version = try_any_cast<const
PythonVersion*>(args.at(0));
+ if (*finished_version != version) {
+ return;
+ }
+ auto* finished_generation =
try_any_cast<std::atomic<int>*>(args.at(1));
+ {
+ std::lock_guard lock(generation_completion->mutex);
+ if (finished_generation !=
generation_completion->generation) {
+ return;
+ }
+ generation_completion->worker_finished = true;
+ }
+ generation_completion->cv.notify_all();
+ });
+ sync_point->enable_processing();
+
+ auto failed_result = mgr._ensure_pool_initialized(version);
+ ASSERT_FALSE(failed_result.has_value());
+
+ // Wait for both parts of the failed generation to finish. Otherwise an
old detached worker can
+ // execute the repaired file and publish its process into the next
generation's pool.
+ {
+ std::unique_lock lock(generation_completion->mutex);
+ ASSERT_TRUE(generation_completion->cv.wait_for(lock,
std::chrono::seconds(5), [&]() {
+ return generation_completion->worker_finished;
+ }));
+ }
+ ASSERT_TRUE(mgr.wait_for_process_pool_initialization_finished_for_test(
+ version, std::chrono::seconds(5)));
- auto start = std::chrono::steady_clock::now();
- auto result = mgr._ensure_pool_initialized(invalid_version);
- auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
- std::chrono::steady_clock::now() - start);
+ // Repair the runtime in place and retry the same version key. Production
can hit this when an
+ // environment or server entry is fixed after a transient initialization
failure.
+ ASSERT_EQ(create_fake_python_with_socket_creation("3.9.16"), python_path);
+ auto recovered_result = mgr._ensure_pool_initialized(version);
+ ASSERT_TRUE(recovered_result.has_value()) <<
recovered_result.error().to_string();
- EXPECT_FALSE(result.has_value());
- EXPECT_LT(elapsed.count(), 500);
+ ProcessPtr process;
+ ASSERT_TRUE(mgr._get_process(version, recovered_result.value(),
&process).ok());
+ ASSERT_NE(process, nullptr);
+ EXPECT_TRUE(process->is_alive());
+
+ mgr.shutdown();
}
TEST_F(PythonServerTest,
EnsurePoolInitializedAfterShutdownReturnsServiceUnavailable) {
@@ -434,13 +453,6 @@ TEST_F(PythonServerTest,
EnsurePoolInitializedAfterShutdownReturnsServiceUnavail
// PythonServerManager::shutdown() - shutdown test
// ============================================================================
-TEST_F(PythonServerTest, ShutdownEmptyManagerDoesNotCrash) {
- PythonServerManager mgr;
-
- // Verify: calling shutdown on empty manager does not crash
- EXPECT_NO_THROW(mgr.shutdown());
-}
-
TEST_F(PythonServerTest, ShutdownCalledMultipleTimesDoesNotCrash) {
PythonServerManager mgr;
@@ -543,32 +555,32 @@ TEST_F(PythonServerTest,
GetClientWithInvalidVersionFails) {
EXPECT_EQ(client, nullptr);
}
-// ============================================================================
-// configuration test
-// ============================================================================
-
-TEST_F(PythonServerTest, MaxPythonProcessNumConfigIsAccessible) {
- // Verify configuration value is accessible and within a valid range
- int max_num = config::max_python_process_num;
- EXPECT_GE(max_num, 0); // 0 means use number of CPU cores
-}
-
// ============================================================================
// destructor test
// ============================================================================
TEST_F(PythonServerTest, DestructorCleansUpResources) {
- // Create and destroy manager to ensure no memory leaks or crashes
+ setup_doris_home();
+ std::string python_path =
create_fake_python_with_socket_creation("3.9.16");
+ config::max_python_process_num = 1;
+
+ ProcessPtr process;
+ std::string socket_path;
{
PythonServerManager mgr;
- // Try some operations (they fail but should not affect destructor)
- PythonVersion invalid_version("3.9.16", "/bad", "/bad");
- ProcessPtr process;
- Status status = mgr.fork(invalid_version, &process);
- EXPECT_FALSE(status.ok());
+ PythonVersion version("3.9.16", test_dir_, python_path);
+ auto pool_result = mgr._ensure_pool_initialized(version);
+ ASSERT_TRUE(pool_result.has_value()) <<
pool_result.error().to_string();
+ ASSERT_TRUE(mgr._get_process(version, pool_result.value(),
&process).ok());
+ ASSERT_NE(process, nullptr);
+ ASSERT_TRUE(process->is_alive());
+ socket_path = process->get_socket_file_path();
+ ASSERT_TRUE(fs::exists(socket_path));
}
- // If we reach here without crashing, destructor works properly
- SUCCEED();
+
+ EXPECT_TRUE(process->is_shutdown());
+ EXPECT_FALSE(process->is_alive());
+ EXPECT_FALSE(fs::exists(socket_path));
}
// ============================================================================
@@ -619,52 +631,6 @@ TEST_F(PythonServerTest, EnsurePoolInitializedSuccess) {
mgr.shutdown();
}
-TEST_F(PythonServerTest,
EnsurePoolInitializedLogsProgressWhileWaitingForSlowProcess) {
- setup_doris_home();
- std::string python_path =
-
create_fake_python_with_delay_and_socket_creation("python3.delayed", "3.9.16",
50);
-
- config::max_python_process_num = 1;
-
- PythonServerManager mgr;
- PythonVersion version("3.9.16", test_dir_, python_path);
-
- auto result = mgr._ensure_pool_initialized(version);
-
- EXPECT_TRUE(result.has_value()) << result.error().to_string();
-
- mgr.shutdown();
-}
-
-TEST_F(PythonServerTest,
EnsurePoolInitializedRetriesAfterInitFailureWithBoundedWait) {
- setup_doris_home();
- std::string python_path =
- create_fake_python_without_socket_creation("python3.no_socket",
"3.9.16");
-
- config::max_python_process_num = 1;
-
- PythonServerManager mgr;
- PythonVersion version("3.9.16", test_dir_, python_path);
-
- auto start = std::chrono::steady_clock::now();
- auto result = mgr._ensure_pool_initialized(version);
- auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
- std::chrono::steady_clock::now() - start);
-
- EXPECT_FALSE(result.has_value());
- EXPECT_LT(elapsed.count(), 2000);
-
- start = std::chrono::steady_clock::now();
- auto retry_result = mgr._ensure_pool_initialized(version);
- elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
- std::chrono::steady_clock::now() - start);
-
- EXPECT_FALSE(retry_result.has_value());
- EXPECT_LT(elapsed.count(), 2000);
-
- mgr.shutdown();
-}
-
TEST_F(PythonServerTest,
EnsurePoolInitializedSucceedsWithOneStuckWorkerAndOneUsableWorker) {
setup_doris_home();
std::string python_path =
@@ -675,23 +641,19 @@ TEST_F(PythonServerTest,
EnsurePoolInitializedSucceedsWithOneStuckWorkerAndOneUs
PythonServerManager mgr;
PythonVersion version("3.9.16", test_dir_, python_path);
- auto start = std::chrono::steady_clock::now();
auto result = mgr._ensure_pool_initialized(version);
- auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
- std::chrono::steady_clock::now() - start);
ASSERT_TRUE(result.has_value()) << result.error().to_string();
- EXPECT_LT(elapsed.count(), 2000);
- EXPECT_TRUE(mgr.process_pool_is_initializing_for_test(version));
ProcessPtr process;
EXPECT_TRUE(mgr._get_process(version, result.value(), &process).ok());
ASSERT_NE(process, nullptr);
EXPECT_TRUE(process->is_alive());
- for (int i = 0; i < 20 &&
!mgr.process_pool_is_initialized_for_test(version); ++i) {
- std::this_thread::sleep_for(std::chrono::milliseconds(100));
- }
+ // Wait on the same condition variable used by the production coordinator
instead of polling
+ // or asserting on a racy intermediate state.
+ EXPECT_TRUE(mgr.wait_for_process_pool_initialized_for_test(version,
+
std::chrono::milliseconds(2000)));
EXPECT_TRUE(mgr.process_pool_is_initialized_for_test(version));
mgr.shutdown();
@@ -708,11 +670,20 @@ TEST_F(PythonServerTest, EnsurePoolInitializedIdempotent)
{
// First initialization
auto result1 = mgr._ensure_pool_initialized(version);
- EXPECT_TRUE(result1.has_value()) << result1.error().to_string();
+ ASSERT_TRUE(result1.has_value()) << result1.error().to_string();
+ auto first_snapshot = mgr.process_pool_snapshot_for_test(version);
+ ASSERT_EQ(first_snapshot.size(), 1);
+ ASSERT_NE(first_snapshot[0], nullptr);
+ pid_t first_pid = first_snapshot[0]->get_child_pid();
- // Second initialization should return immediately (version already
initialized)
+ // Re-initialization must reuse both the versioned pool and its live
process.
auto result2 = mgr._ensure_pool_initialized(version);
- EXPECT_TRUE(result2.has_value()) << result2.error().to_string();
+ ASSERT_TRUE(result2.has_value()) << result2.error().to_string();
+ EXPECT_EQ(result2.value(), result1.value());
+ auto second_snapshot = mgr.process_pool_snapshot_for_test(version);
+ ASSERT_EQ(second_snapshot.size(), 1);
+ EXPECT_EQ(second_snapshot[0], first_snapshot[0]);
+ EXPECT_EQ(second_snapshot[0]->get_child_pid(), first_pid);
mgr.shutdown();
}
@@ -811,30 +782,22 @@ TEST_F(PythonServerTest,
GetProcessSkipsDeadProcessWhenAliveProcessExists) {
mgr.shutdown();
}
-TEST_F(PythonServerTest, GetProcessLoadBalancing) {
- setup_doris_home();
- std::string python_path =
create_fake_python_with_socket_creation("3.9.16");
-
- // Create a pool with 2 processes
- config::max_python_process_num = 2;
-
+TEST_F(PythonServerTest, GetProcessSelectsLeastSharedProcess) {
PythonServerManager mgr;
- PythonVersion version("3.9.16", test_dir_, python_path);
+ PythonVersion version("3.9.16", test_dir_, test_dir_ + "/unused_python");
+ mgr.set_process_pool_for_test(version, {create_sleep_process(),
create_sleep_process()});
auto init_result = mgr._ensure_pool_initialized(version);
EXPECT_TRUE(init_result.has_value()) << init_result.error().to_string();
- // Get multiple processes to verify load balancing
- ProcessPtr p1, p2, p3, p4;
- EXPECT_TRUE(mgr._get_process(version, init_result.value(), &p1).ok());
- EXPECT_TRUE(mgr._get_process(version, init_result.value(), &p2).ok());
- EXPECT_TRUE(mgr._get_process(version, init_result.value(), &p3).ok());
- EXPECT_TRUE(mgr._get_process(version, init_result.value(), &p4).ok());
-
- // With 2 processes, load balancing distributes requests across different
processes
- // p1 and p2 may be same or different processes
- EXPECT_NE(p1, nullptr);
- EXPECT_NE(p2, nullptr);
+ // Holding p1 increases its shared ownership count, so the next client
must select the other
+ // live process. An implementation that always returns pool[0] fails this
assertion.
+ ProcessPtr p1, p2;
+ ASSERT_TRUE(mgr._get_process(version, init_result.value(), &p1).ok());
+ ASSERT_NE(p1, nullptr);
+ ASSERT_TRUE(mgr._get_process(version, init_result.value(), &p2).ok());
+ ASSERT_NE(p2, nullptr);
+ EXPECT_NE(p1->get_child_pid(), p2->get_child_pid());
mgr.shutdown();
}
@@ -924,33 +887,82 @@ TEST_F(PythonServerTest,
EnsurePoolInitializedForDifferentVersionsDoesNotShareVe
config::max_python_process_num = 1;
- std::string python39_path =
- create_fake_python_with_delay_and_socket_creation("python3.9",
"3.9.16", 50);
- std::string python310_path =
- create_fake_python_with_delay_and_socket_creation("python3.10",
"3.10.0", 50);
+ std::string python39_path =
create_fake_python_with_socket_creation("3.9.16");
+ std::string python310_path = test_dir_ + "/bin/python3.10";
+ ASSERT_TRUE(fs::copy_file(python39_path, python310_path));
+ fs::permissions(python310_path, fs::perms::owner_all);
PythonServerManager mgr;
PythonVersion version39("3.9.16", test_dir_, python39_path);
PythonVersion version310("3.10.0", test_dir_, python310_path);
- auto start = std::chrono::steady_clock::now();
+ struct ForkBarrier {
+ std::mutex mutex;
+ std::condition_variable cv;
+ int entries = 0;
+ bool released = false;
+ };
+ auto fork_barrier = std::make_shared<ForkBarrier>();
+ auto* sync_point = SyncPoint::get_instance();
+ Defer clear_sync_point {[fork_barrier, sync_point]() {
+ {
+ std::lock_guard lock(fork_barrier->mutex);
+ fork_barrier->released = true;
+ }
+ fork_barrier->cv.notify_all();
+ sync_point->disable_processing();
+
sync_point->clear_call_back("PythonServerManager::fork:before_process_start");
+ sync_point->clear_call_back(
+
"PythonServerManager::_ensure_pool_initialized:process_pool_init_timeout");
+ }};
+ // Keep callers inside initialization longer than the fork barrier. A
manager-wide lock must
+ // therefore fail the barrier instead of serializing through the short
BE_TEST pool timeout.
+ sync_point->set_call_back(
+
"PythonServerManager::_ensure_pool_initialized:process_pool_init_timeout",
+ [](auto&& args) {
+ auto* timeout =
try_any_cast<std::chrono::milliseconds*>(args.front());
+ *timeout = std::chrono::seconds(10);
+ });
+ sync_point->set_call_back(
+ "PythonServerManager::fork:before_process_start",
+ [fork_barrier, version39, version310](auto&& args) {
+ if (args.empty()) {
+ return;
+ }
+ const auto* entering_version = try_any_cast<const
PythonVersion*>(args.front());
+ if (*entering_version != version39 && *entering_version !=
version310) {
+ return;
+ }
+ std::unique_lock lock(fork_barrier->mutex);
+ ++fork_barrier->entries;
+ fork_barrier->cv.notify_all();
+ fork_barrier->cv.wait(lock, [&]() { return
fork_barrier->released; });
+ });
+ sync_point->enable_processing();
+
auto future39 = std::async(std::launch::async,
[&]() { return
mgr._ensure_pool_initialized(version39); });
auto future310 = std::async(std::launch::async,
[&]() { return
mgr._ensure_pool_initialized(version310); });
- auto result39 = future39.get();
- auto result310 = future310.get();
- auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
- std::chrono::steady_clock::now() - start);
+ bool both_forks_entered = false;
+ {
+ std::unique_lock lock(fork_barrier->mutex);
+ both_forks_entered = fork_barrier->cv.wait_for(
+ lock, std::chrono::seconds(5), [&]() { return
fork_barrier->entries >= 2; });
+ fork_barrier->released = true;
+ }
+ fork_barrier->cv.notify_all();
- EXPECT_TRUE(result39.has_value()) << result39.error().to_string();
- EXPECT_TRUE(result310.has_value()) << result310.error().to_string();
- // Keep the assertion loose for ASAN/CI scheduling while still catching
full init-timeout
- // serialization between versions.
- EXPECT_LT(elapsed.count(), 2000);
+ // The assertion below is the concurrency evidence. Pool completion also
has a separate short
+ // BE_TEST timeout, so its status must not turn host scheduling after this
barrier into noise.
+ static_cast<void>(future39.get());
+ static_cast<void>(future310.get());
mgr.shutdown();
+
+ EXPECT_TRUE(both_forks_entered)
+ << "Both version pools must enter fork before either process
starts";
}
// ============================================================================
diff --git a/be/test/udf/python/python_udf_meta_test.cpp
b/be/test/udf/python/python_udf_meta_test.cpp
index 43085430510..66067b37f35 100644
--- a/be/test/udf/python/python_udf_meta_test.cpp
+++ b/be/test/udf/python/python_udf_meta_test.cpp
@@ -17,6 +17,9 @@
#include "udf/python/python_udf_meta.h"
+#include <arrow/io/memory.h>
+#include <arrow/ipc/reader.h>
+#include <arrow/util/base64.h>
#include <gtest/gtest.h>
#include <rapidjson/document.h>
@@ -29,6 +32,14 @@
namespace doris {
+static arrow::Result<std::shared_ptr<arrow::Schema>> decode_arrow_schema(
+ const std::string& encoded_schema) {
+ auto buffer =
arrow::Buffer::FromString(arrow::util::base64_decode(encoded_schema));
+ arrow::io::BufferReader reader(buffer);
+ arrow::ipc::DictionaryMemo dictionary_memo;
+ return arrow::ipc::ReadSchema(&reader, &dictionary_memo);
+}
+
class PythonUDFMetaTest : public ::testing::Test {
protected:
void SetUp() override {
@@ -46,23 +57,6 @@ protected:
DataTypePtr nullable_double_;
};
-// ============================================================================
-// PythonUDFMeta construction tests
-// ============================================================================
-
-TEST_F(PythonUDFMetaTest, DefaultConstruction) {
- PythonUDFMeta meta;
- EXPECT_TRUE(meta.name.empty());
- EXPECT_TRUE(meta.symbol.empty());
- EXPECT_TRUE(meta.location.empty());
- EXPECT_TRUE(meta.checksum.empty());
- EXPECT_TRUE(meta.runtime_version.empty());
- EXPECT_TRUE(meta.inline_code.empty());
- EXPECT_FALSE(meta.always_nullable);
- EXPECT_TRUE(meta.input_types.empty());
- EXPECT_EQ(meta.return_type, nullptr);
-}
-
// ============================================================================
// PythonUDFMeta check() tests
// ============================================================================
@@ -254,77 +248,6 @@ TEST_F(PythonUDFMetaTest, CheckWhitespaceOnlyName) {
EXPECT_FALSE(status.ok());
}
-// ============================================================================
-// PythonUDFMeta to_string() tests
-// ============================================================================
-
-TEST_F(PythonUDFMetaTest, ToStringContainsAllFields) {
- PythonUDFMeta meta;
- meta.name = "my_udf";
- meta.symbol = "udf_func";
- meta.location = "/path/to/udf.py";
- meta.runtime_version = "3.10.5";
- meta.always_nullable = true;
- meta.inline_code = "def udf_func(x): return x";
- meta.input_types = {nullable_int32_, nullable_string_};
- meta.return_type = nullable_double_;
-
- std::string str = meta.to_string();
- EXPECT_TRUE(str.find("my_udf") != std::string::npos);
- EXPECT_TRUE(str.find("udf_func") != std::string::npos);
- EXPECT_TRUE(str.find("/path/to/udf.py") != std::string::npos);
- EXPECT_TRUE(str.find("3.10.5") != std::string::npos);
-}
-
-TEST_F(PythonUDFMetaTest, ToStringMultipleInputTypes) {
- PythonUDFMeta meta;
- meta.name = "multi_arg_udf";
- meta.symbol = "func";
- meta.runtime_version = "3.9.16";
- meta.input_types = {nullable_int32_, nullable_string_, nullable_double_};
- meta.return_type = nullable_int32_;
-
- std::string str = meta.to_string();
- // Should contain input_types section
- EXPECT_TRUE(str.find("input_types") != std::string::npos);
-}
-
-// ============================================================================
-// PythonUDFMeta equality tests
-// ============================================================================
-
-TEST_F(PythonUDFMetaTest, EqualityById) {
- PythonUDFMeta meta1;
- meta1.id = 100;
- meta1.name = "udf1";
-
- PythonUDFMeta meta2;
- meta2.id = 100;
- meta2.name = "different_name";
-
- PythonUDFMeta meta3;
- meta3.id = 200;
- meta3.name = "udf1";
-
- EXPECT_EQ(meta1, meta2); // Same ID
- EXPECT_FALSE(meta1 == meta3); // Different ID
-}
-
-TEST_F(PythonUDFMetaTest, HashById) {
- PythonUDFMeta meta1;
- meta1.id = 100;
-
- PythonUDFMeta meta2;
- meta2.id = 100;
-
- PythonUDFMeta meta3;
- meta3.id = 200;
-
- std::hash<PythonUDFMeta> hasher;
- EXPECT_EQ(hasher(meta1), hasher(meta2));
- EXPECT_NE(hasher(meta1), hasher(meta3));
-}
-
// ============================================================================
// PythonUDFMeta serialize_to_json() tests
// ============================================================================
@@ -379,6 +302,22 @@ TEST_F(PythonUDFMetaTest, SerializeToJsonBasic) {
EXPECT_TRUE(doc.HasMember("inline_code"));
EXPECT_TRUE(doc.HasMember("input_types"));
EXPECT_TRUE(doc.HasMember("return_type"));
+
+ EXPECT_EQ(arrow::util::base64_decode(doc["inline_code"].GetString()),
meta.inline_code);
+
+ auto input_schema_result =
decode_arrow_schema(doc["input_types"].GetString());
+ ASSERT_TRUE(input_schema_result.ok()) <<
input_schema_result.status().ToString();
+ auto input_schema = *input_schema_result;
+ ASSERT_EQ(input_schema->num_fields(), 1);
+ EXPECT_TRUE(input_schema->field(0)->type()->Equals(arrow::int32()));
+ EXPECT_TRUE(input_schema->field(0)->nullable());
+
+ auto return_schema_result =
decode_arrow_schema(doc["return_type"].GetString());
+ ASSERT_TRUE(return_schema_result.ok()) <<
return_schema_result.status().ToString();
+ auto return_schema = *return_schema_result;
+ ASSERT_EQ(return_schema->num_fields(), 1);
+ EXPECT_TRUE(return_schema->field(0)->type()->Equals(arrow::int32()));
+ EXPECT_TRUE(return_schema->field(0)->nullable());
}
TEST_F(PythonUDFMetaTest, SerializeToJsonDifferentClientTypes) {
@@ -430,7 +369,16 @@ TEST_F(PythonUDFMetaTest,
SerializeToJsonMultipleInputTypes) {
rapidjson::Document doc;
doc.Parse(json_str.c_str());
EXPECT_FALSE(doc.HasParseError());
- EXPECT_TRUE(doc.HasMember("input_types"));
+ auto input_schema_result =
decode_arrow_schema(doc["input_types"].GetString());
+ ASSERT_TRUE(input_schema_result.ok()) <<
input_schema_result.status().ToString();
+ auto input_schema = *input_schema_result;
+ ASSERT_EQ(input_schema->num_fields(), 3);
+ EXPECT_TRUE(input_schema->field(0)->type()->Equals(arrow::int32()));
+ EXPECT_TRUE(input_schema->field(1)->type()->Equals(arrow::utf8()));
+ EXPECT_TRUE(input_schema->field(2)->type()->Equals(arrow::float64()));
+ EXPECT_TRUE(input_schema->field(0)->nullable());
+ EXPECT_TRUE(input_schema->field(1)->nullable());
+ EXPECT_TRUE(input_schema->field(2)->nullable());
}
TEST_F(PythonUDFMetaTest, SerializeToJsonEmptyInputTypesForUdf) {
@@ -450,8 +398,9 @@ TEST_F(PythonUDFMetaTest,
SerializeToJsonEmptyInputTypesForUdf) {
rapidjson::Document doc;
doc.Parse(json_str.c_str());
EXPECT_FALSE(doc.HasParseError());
- EXPECT_TRUE(doc.HasMember("input_types"));
- EXPECT_FALSE(std::string(doc["input_types"].GetString()).empty());
+ auto input_schema_result =
decode_arrow_schema(doc["input_types"].GetString());
+ ASSERT_TRUE(input_schema_result.ok()) <<
input_schema_result.status().ToString();
+ EXPECT_EQ((*input_schema_result)->num_fields(), 0);
}
// ============================================================================
@@ -469,6 +418,10 @@ TEST_F(PythonUDFMetaTest, ConvertTypesToSchemaBasic) {
EXPECT_EQ(schema->num_fields(), 2);
EXPECT_EQ(schema->field(0)->name(), "arg0");
EXPECT_EQ(schema->field(1)->name(), "arg1");
+ EXPECT_TRUE(schema->field(0)->type()->Equals(arrow::int32()));
+ EXPECT_TRUE(schema->field(1)->type()->Equals(arrow::utf8()));
+ EXPECT_TRUE(schema->field(0)->nullable());
+ EXPECT_TRUE(schema->field(1)->nullable());
}
TEST_F(PythonUDFMetaTest, ConvertTypesToSchemaSingleType) {
@@ -480,6 +433,9 @@ TEST_F(PythonUDFMetaTest, ConvertTypesToSchemaSingleType) {
EXPECT_TRUE(status.ok()) << status.to_string();
EXPECT_NE(schema, nullptr);
EXPECT_EQ(schema->num_fields(), 1);
+ EXPECT_EQ(schema->field(0)->name(), "arg0");
+ EXPECT_TRUE(schema->field(0)->type()->Equals(arrow::float64()));
+ EXPECT_TRUE(schema->field(0)->nullable());
}
TEST_F(PythonUDFMetaTest, ConvertTypesToSchemaEmpty) {
@@ -506,6 +462,12 @@ TEST_F(PythonUDFMetaTest, SerializeArrowSchema) {
EXPECT_TRUE(status.ok()) << status.to_string();
EXPECT_NE(buffer, nullptr);
EXPECT_GT(buffer->size(), 0);
+
+ arrow::io::BufferReader reader(buffer);
+ arrow::ipc::DictionaryMemo dictionary_memo;
+ auto decoded_schema_result = arrow::ipc::ReadSchema(&reader,
&dictionary_memo);
+ ASSERT_TRUE(decoded_schema_result.ok()) <<
decoded_schema_result.status().ToString();
+ EXPECT_TRUE((*decoded_schema_result)->Equals(*schema));
}
} // namespace doris
diff --git a/be/test/udf/python/python_udf_runtime_test.cpp
b/be/test/udf/python/python_udf_runtime_test.cpp
index 0f570c06810..4b07b0b7fb8 100644
--- a/be/test/udf/python/python_udf_runtime_test.cpp
+++ b/be/test/udf/python/python_udf_runtime_test.cpp
@@ -25,8 +25,11 @@
#include <boost/process.hpp>
#include <filesystem>
+#include <optional>
#include <string>
+#include "util/defer_op.h"
+
namespace doris {
namespace fs = std::filesystem;
@@ -69,41 +72,18 @@ protected:
// Helper function tests
// ============================================================================
-TEST_F(PythonUDFRuntimeTest, GetBaseUnixSocketPath) {
- std::string path = get_base_unix_socket_path();
- EXPECT_TRUE(path.find("grpc+unix://") != std::string::npos);
- EXPECT_TRUE(path.find("/tmp/doris_python_udf") != std::string::npos);
-}
-
-TEST_F(PythonUDFRuntimeTest, GetUnixSocketPath) {
- pid_t test_pid = 12345;
- std::string path = get_unix_socket_path(test_pid);
- EXPECT_TRUE(path.find("grpc+unix://") != std::string::npos);
- EXPECT_TRUE(path.find("12345") != std::string::npos);
- EXPECT_TRUE(path.find(".sock") != std::string::npos);
-}
-
-TEST_F(PythonUDFRuntimeTest, GetUnixSocketFilePathFormat) {
- pid_t test_pid = 99999;
- std::string path = get_unix_socket_file_path(test_pid);
- EXPECT_TRUE(path.find("/tmp/doris_python_udf") != std::string::npos);
- EXPECT_TRUE(path.find("99999") != std::string::npos);
- EXPECT_TRUE(path.find(".sock") != std::string::npos);
- // Should NOT have grpc+unix:// prefix
- EXPECT_TRUE(path.find("grpc+unix://") == std::string::npos);
-}
-
-TEST_F(PythonUDFRuntimeTest, GetUnixSocketPathDifferentPids) {
- std::string path1 = get_unix_socket_path(1000);
- std::string path2 = get_unix_socket_path(2000);
- EXPECT_NE(path1, path2);
- EXPECT_TRUE(path1.find("1000") != std::string::npos);
- EXPECT_TRUE(path2.find("2000") != std::string::npos);
+TEST_F(PythonUDFRuntimeTest, UnixSocketPathsMatchFlightContract) {
+ EXPECT_EQ(get_base_unix_socket_path(),
"grpc+unix:///tmp/doris_python_udf");
+ EXPECT_EQ(get_unix_socket_path(12345),
"grpc+unix:///tmp/doris_python_udf_12345.sock");
+ EXPECT_EQ(get_unix_socket_file_path(12345),
"/tmp/doris_python_udf_12345.sock");
}
TEST_F(PythonUDFRuntimeTest, GetFightServerPath) {
// Save original DORIS_HOME
- const char* original_doris_home = std::getenv("DORIS_HOME");
+ std::optional<std::string> original_doris_home;
+ if (const char* doris_home = std::getenv("DORIS_HOME")) {
+ original_doris_home = doris_home;
+ }
// Set test DORIS_HOME
setenv("DORIS_HOME", "/test/doris/home", 1);
@@ -113,26 +93,12 @@ TEST_F(PythonUDFRuntimeTest, GetFightServerPath) {
// Restore original DORIS_HOME
if (original_doris_home) {
- setenv("DORIS_HOME", original_doris_home, 1);
+ setenv("DORIS_HOME", original_doris_home->c_str(), 1);
} else {
unsetenv("DORIS_HOME");
}
}
-// ============================================================================
-// PythonUDFProcess tests (without actually spawning a process)
-// ============================================================================
-
-// Note: Most PythonUDFProcess tests require actually spawning Python
processes,
-// which is environment-dependent. Here we test what we can without real
processes.
-
-TEST_F(PythonUDFRuntimeTest, ProcessPtrIsSharedPtr) {
- // Verify ProcessPtr is a shared_ptr
- ProcessPtr ptr = nullptr;
- EXPECT_EQ(ptr, nullptr);
- EXPECT_FALSE(ptr);
-}
-
TEST_F(PythonUDFRuntimeTest, WaitChildExitReturnsExitedForExitedChild) {
bp::ipstream output;
bp::child child("/bin/bash", "-c", "exit 7", bp::std_out > output);
@@ -203,52 +169,6 @@ TEST_F(PythonUDFRuntimeTest,
BackgroundReaperReapsQueuedChild) {
EXPECT_EQ(result, PythonUDFProcess::ChildExitWaitResult::ALREADY_REAPED);
}
-// Test socket file path generation for various PIDs
-TEST_F(PythonUDFRuntimeTest, SocketPathGenerationEdgeCases) {
- // Minimum PID
- std::string path1 = get_unix_socket_file_path(1);
- EXPECT_TRUE(path1.find("_1.sock") != std::string::npos);
-
- // Large PID
- std::string path2 = get_unix_socket_file_path(999999);
- EXPECT_TRUE(path2.find("_999999.sock") != std::string::npos);
-}
-
-TEST_F(PythonUDFRuntimeTest, SocketPathConsistency) {
- pid_t pid = 54321;
- // Multiple calls should return same result
- std::string path1 = get_unix_socket_path(pid);
- std::string path2 = get_unix_socket_path(pid);
- EXPECT_EQ(path1, path2);
-
- std::string file_path1 = get_unix_socket_file_path(pid);
- std::string file_path2 = get_unix_socket_file_path(pid);
- EXPECT_EQ(file_path1, file_path2);
-}
-
-TEST_F(PythonUDFRuntimeTest, UnixSocketPathRelationship) {
- pid_t pid = 12345;
- std::string socket_path = get_unix_socket_path(pid);
- std::string file_path = get_unix_socket_file_path(pid);
-
- // Socket path should contain the prefix + file path structure
- EXPECT_TRUE(socket_path.find(UNIX_SOCKET_PREFIX) == 0);
-
- // File path should not contain the prefix
- EXPECT_TRUE(file_path.find(UNIX_SOCKET_PREFIX) == std::string::npos);
-}
-
-// ============================================================================
-// Socket path security tests
-// ============================================================================
-
-TEST_F(PythonUDFRuntimeTest, SocketPathInTmpDirectory) {
- // All socket files should be in /tmp to avoid path length issues
- pid_t pid = 12345;
- std::string file_path = get_unix_socket_file_path(pid);
- EXPECT_TRUE(file_path.find("/tmp/") == 0);
-}
-
TEST_F(PythonUDFRuntimeTest, SocketPathLength) {
// Unix socket paths have a maximum length (usually 107 chars)
// Verify generated paths are within reasonable limits
@@ -259,18 +179,6 @@ TEST_F(PythonUDFRuntimeTest, SocketPathLength) {
EXPECT_LT(path.length(), 100);
}
-// ============================================================================
-// Flight server path tests
-// ============================================================================
-
-TEST_F(PythonUDFRuntimeTest, FlightServerPathTemplate) {
- // Verify template includes necessary components
- std::string tmpl = FLIGHT_SERVER_PATH_TEMPLATE;
- EXPECT_TRUE(tmpl.find("{}") != std::string::npos); // Has placeholder
- EXPECT_TRUE(tmpl.find("plugins") != std::string::npos);
- EXPECT_TRUE(tmpl.find("python_udf") != std::string::npos);
-}
-
// ============================================================================
// PythonUDFProcess shutdown() tests
// ============================================================================
@@ -375,126 +283,43 @@ TEST_F(PythonUDFRuntimeTest,
RemoveUnixSocketExistingFile) {
EXPECT_FALSE(fs::exists(socket_path));
}
-TEST_F(PythonUDFRuntimeTest, RemoveUnixSocketNonExistent) {
+TEST_F(PythonUDFRuntimeTest, ShutdownPreservesUnexpectedSocketDirectory) {
bp::ipstream output;
bp::child child("/bin/sleep", "60", bp::std_out > output);
+ ASSERT_TRUE(child.valid());
+ ASSERT_TRUE(child.running());
- pid_t child_pid = child.id();
PythonUDFProcess process(std::move(child), std::move(output));
+ const std::string socket_path = process.get_socket_file_path();
+ fs::remove_all(socket_path);
+ ASSERT_TRUE(fs::create_directory(socket_path));
+ Defer cleanup {[&]() { fs::remove_all(socket_path); }};
- // Don't create socket file - it doesn't exist
- std::string socket_path = get_unix_socket_file_path(child_pid);
- ASSERT_FALSE(fs::exists(socket_path));
-
- // Shutdown should not crash even if socket doesn't exist (ENOENT case)
+ // A directory at the derived socket path makes unlink() fail with EISDIR.
Shutdown must still
+ // finish, and it must not recursively delete an unexpected filesystem
object.
process.shutdown();
- EXPECT_TRUE(process.is_shutdown());
-}
-
-TEST_F(PythonUDFRuntimeTest, RemoveUnixSocketIsDirectory) {
- bp::ipstream output;
- bp::child child("/bin/sleep", "60", bp::std_out > output);
-
- pid_t child_pid = child.id();
- // Create a directory at the socket path location (instead of a file)
- // This will cause unlink() to fail with EISDIR
- std::string socket_path = get_unix_socket_file_path(child_pid);
- fs::create_directories(socket_path);
- ASSERT_TRUE(fs::is_directory(socket_path));
-
- PythonUDFProcess process(std::move(child), std::move(output));
-
- // Shutdown should handle EISDIR error gracefully (logs warning but
doesn't crash)
- process.shutdown();
EXPECT_TRUE(process.is_shutdown());
-
- // Cleanup - remove the directory we created
- fs::remove_all(socket_path);
-}
-
-// ============================================================================
-// PythonUDFProcess to_string() tests
-// ============================================================================
-
-TEST_F(PythonUDFRuntimeTest, ToStringFormat) {
- bp::ipstream output;
- bp::child child("/bin/sleep", "60", bp::std_out > output);
-
- pid_t child_pid = child.id();
- PythonUDFProcess process(std::move(child), std::move(output));
-
- std::string str = process.to_string();
-
- // Verify to_string contains expected fields
- EXPECT_TRUE(str.find("PythonUDFProcess") != std::string::npos);
- EXPECT_TRUE(str.find("child_pid") != std::string::npos);
- EXPECT_TRUE(str.find(std::to_string(child_pid)) != std::string::npos);
- EXPECT_TRUE(str.find("uri") != std::string::npos);
- EXPECT_TRUE(str.find("unix_socket_file_path") != std::string::npos);
- EXPECT_TRUE(str.find("is_shutdown") != std::string::npos);
-
- process.shutdown();
+ EXPECT_TRUE(fs::is_directory(socket_path));
}
// ============================================================================
// PythonUDFProcess getter tests
// ============================================================================
-TEST_F(PythonUDFRuntimeTest, GetUri) {
- bp::ipstream output;
- bp::child child("/bin/sleep", "60", bp::std_out > output);
-
- pid_t child_pid = child.id();
- PythonUDFProcess process(std::move(child), std::move(output));
-
- std::string uri = process.get_uri();
-
- // URI should contain the grpc+unix prefix and pid
- EXPECT_TRUE(uri.find("grpc+unix://") != std::string::npos);
- EXPECT_TRUE(uri.find(std::to_string(child_pid)) != std::string::npos);
- EXPECT_TRUE(uri.find(".sock") != std::string::npos);
-
- process.shutdown();
-}
-
-TEST_F(PythonUDFRuntimeTest, GetSocketFilePath) {
+TEST_F(PythonUDFRuntimeTest, ProcessUsesPidDerivedSocketPaths) {
bp::ipstream output;
bp::child child("/bin/sleep", "60", bp::std_out > output);
pid_t child_pid = child.id();
PythonUDFProcess process(std::move(child), std::move(output));
- const std::string& path = process.get_socket_file_path();
-
- // File path should NOT have grpc+unix prefix
- EXPECT_TRUE(path.find("grpc+unix://") == std::string::npos);
- EXPECT_TRUE(path.find(std::to_string(child_pid)) != std::string::npos);
- EXPECT_TRUE(path.find(".sock") != std::string::npos);
- EXPECT_TRUE(path.find("/tmp/") == 0);
+ EXPECT_EQ(process.get_uri(), get_unix_socket_path(child_pid));
+ EXPECT_EQ(process.get_socket_file_path(),
get_unix_socket_file_path(child_pid));
process.shutdown();
}
-// ============================================================================
-// PythonUDFProcess equality tests
-// ============================================================================
-
-TEST_F(PythonUDFRuntimeTest, ProcessEquality) {
- bp::ipstream output1, output2;
- bp::child child1("/bin/sleep", "60", bp::std_out > output1);
- bp::child child2("/bin/sleep", "60", bp::std_out > output2);
-
- PythonUDFProcess process1(std::move(child1), std::move(output1));
- PythonUDFProcess process2(std::move(child2), std::move(output2));
-
- // Different processes should not be equal (different PIDs)
- EXPECT_NE(process1, process2);
-
- process1.shutdown();
- process2.shutdown();
-}
-
// ============================================================================
// PythonUDFProcess destructor tests
// ============================================================================
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
