diff --git a/python/pyarrow/tests/test_serialization.py b/python/pyarrow/tests/test_serialization.py index 0917172d22c..feccebbde36 100644 --- a/python/pyarrow/tests/test_serialization.py +++ b/python/pyarrow/tests/test_serialization.py @@ -580,6 +580,22 @@ def deserialize_regex(serialized, q): p.join() +def _get_modified_env_with_pythonpath(): + # Prepend pyarrow root directory to PYTHONPATH + env = os.environ.copy() + existing_pythonpath = env.get('PYTHONPATH', '') + if sys.platform == 'win32': + sep = ';' + else: + sep = ':' + + module_path = os.path.abspath( + os.path.dirname(os.path.dirname(pa.__file__))) + + env['PYTHONPATH'] = sep.join((module_path, existing_pythonpath)) + return env + + def test_deserialize_buffer_in_different_process(): import tempfile import subprocess @@ -589,9 +605,12 @@ def test_deserialize_buffer_in_different_process(): f.write(b.to_pybytes()) f.close() + subprocess_env = _get_modified_env_with_pythonpath() + dir_path = os.path.dirname(os.path.realpath(__file__)) python_file = os.path.join(dir_path, 'deserialize_buffer.py') - subprocess.check_call([sys.executable, python_file, f.name]) + subprocess.check_call([sys.executable, python_file, f.name], + env=subprocess_env) def test_set_pickle():