Google Git Google Git
Sign in
apache/spark/refs/heads/master/./python/pyspark/sql/tests/test_inprocess_udf.py
blob: b12f24d583d54f0f0993586dbec5b751c94628cb [file]
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""End-to-end tests for in-process Python UDFs.
Run with python/run-tests like other SQL tests. JEP paths are discovered from the
selected Python environment before the Spark JVM starts. ARROW_C_DATA_JAR must
point to the provided Arrow CDI JAR. Set INPROCESS_TESTS=1
to require the suite (missing dependencies then fail), or 0 to disable it.
Otherwise, the suite runs when JEP, PyArrow and the CDI JAR are available.
"""
import os
import shutil
import tempfile
import unittest
import zipfile
from importlib.util import find_spec
from pathlib import Path
from unittest.mock import patch
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.testing.utils import eventually, have_pyarrow
_jep_spec = find_spec("jep")
_cdi_jar = os.environ.get("ARROW_C_DATA_JAR")
_test_mode = os.environ.get("INPROCESS_TESTS")
_run_inprocess = _test_mode == "1" or (
_test_mode != "0"
and _jep_spec is not None
and have_pyarrow
and _cdi_jar is not None
and Path(_cdi_jar).is_file()
)
@unittest.skipUnless(_run_inprocess, "Requires JEP, PyArrow and ARROW_C_DATA_JAR")
class InProcessUDFTests(ReusedSQLTestCase):
"""
End-to-end tests for @inprocess_udf that require jep + CPython + PyArrow.
The plugin initializes JEP before any task starts. Calls from task threads and
shutdown must use the same dedicated interpreter thread.
"""
@classmethod
def master(cls):
return "local[2]"
@classmethod
def conf(cls):
return (
super()
.conf()
.set("spark.task.cpus", "0.5")
.set("spark.driver.extraClassPath", os.pathsep.join([str(cls.jep_jar), cls.cdi_jar]))
.set("spark.driver.extraLibraryPath", str(cls.jep_dir))
.set(
"spark.python.inProcess.sitePackages",
",".join([cls.site_packages, str(cls.jep_dir.parent)]),
)
.set("spark.plugins", "org.apache.spark.sql.execution.python.InProcessPythonPlugin")
)
@classmethod
def setUpClass(cls):
if _jep_spec is None:
raise RuntimeError("INPROCESS_TESTS=1 requires JEP in the selected Python environment")
# Do not import jep: it can only be imported by an embedded interpreter.
cls.jep_dir = Path(_jep_spec.origin).parent
jars = list(cls.jep_dir.glob("jep-*.jar"))
if len(jars) != 1:
raise RuntimeError(f"Expected one JEP JAR in {cls.jep_dir}, found {len(jars)}")
cls.jep_jar = jars[0]
if not _cdi_jar or not Path(_cdi_jar).is_file():
raise RuntimeError("Set ARROW_C_DATA_JAR to the provided Arrow CDI JAR")
cls.cdi_jar = str(Path(_cdi_jar).resolve())
cls.site_packages = tempfile.mkdtemp()
helper_dir = os.path.join(cls.site_packages, "extra")
system_dir = os.path.join(cls.site_packages, "system")
os.mkdir(helper_dir)
os.mkdir(system_dir)
with open(os.path.join(system_dir, "_inprocess_process_helper.py"), "w") as f:
f.write("MAGIC = -1\n")
with open(os.path.join(helper_dir, "_inprocess_test_helper.py"), "w") as f:
f.write("MAGIC = 99\n")
with open(os.path.join(cls.site_packages, "helper.pth"), "w") as f:
f.write("extra\n")
for name in ["spire", "redis"]:
package = Path(cls.site_packages) / name
package.mkdir()
(package / "__init__.py").write_text("PYTHON_PACKAGE = True\n")
shadow = os.path.join(cls.site_packages, "pyspark")
os.mkdir(shadow)
with open(os.path.join(shadow, "__init__.py"), "w") as f:
f.write("raise RuntimeError('site-packages must not shadow Spark PySpark')\n")
try:
# CI extracts compiled targets without building pyspark.zip. Keep the source
# tree available, while still requiring the configured paths to provide JEP.
cls.python_source = Path(__file__).resolve().parents[3]
archive = cls.python_source / "lib" / "pyspark.zip"
if archive.is_file():
with zipfile.ZipFile(archive) as packaged:
for module in (cls.python_source / "pyspark" / "inprocess").glob("*.py"):
name = module.relative_to(cls.python_source).as_posix()
if (
name not in packaged.namelist()
or packaged.read(name) != module.read_bytes()
):
raise RuntimeError("Rebuild or remove stale python/lib/pyspark.zip")
python_path = os.pathsep.join([str(cls.python_source), system_dir])
with patch.dict(os.environ, {"PYTHONPATH": python_path}):
super().setUpClass()
except Exception:
shutil.rmtree(cls.site_packages)
raise
@classmethod
def tearDownClass(cls):
try:
super().tearDownClass()
finally:
shutil.rmtree(cls.site_packages)
def test_driver_defined_udt_return_type(self):
import sys
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType, UserDefinedType
class DriverUDT(UserDefinedType):
@classmethod
def sqlType(cls):
return LongType()
@classmethod
def module(cls):
return "__main__"
def serialize(self, value):
return value
def deserialize(self, value):
return value
DriverUDT.__module__ = "__main__"
with patch.object(sys.modules["__main__"], "DriverUDT", DriverUDT, create=True):
identity = inprocess_udf(DriverUDT())(lambda x: x)
result = self.spark.range(3).select(identity("id"))
self.assertEqual([r[0] for r in result.collect()], [0, 1, 2])
def test_unsupported_ddl_return_type_fails_on_driver(self):
from pyspark.errors import PySparkNotImplementedError
from pyspark.inprocess import inprocess_udf
wrapper = inprocess_udf("interval year to month")(lambda x: x)
with patch("pyspark.inprocess.udf._serialize_udf") as serialize:
with self.assertRaises(PySparkNotImplementedError):
wrapper("id")
serialize.assert_not_called()
self.assertIsNone(wrapper._jfunc)
def test_worker_environment_is_rejected(self):
from pyspark.inprocess import inprocess_udf
identity = inprocess_udf("long")(lambda x: x)
column = identity("id")
with self.sql_conf({"spark.pythonWorkerEnv.INPROCESS_TEST_VALUE": "value"}):
# Columns are session-agnostic; validate using the query's session at execution.
identity("id")
with self.assertRaisesRegex(Exception, "UNSUPPORTED_IN_PROCESS_PYTHON_UDF"):
self.spark.range(1).select(column).collect()
def test_profiler_and_memory_settings_are_rejected(self):
from pyspark.inprocess import inprocess_udf
identity = inprocess_udf("long")(lambda x: x)
column = identity("id")
with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}):
# Columns are session-agnostic; validate using the query's session at execution.
identity("id")
with self.assertRaisesRegex(Exception, "UNSUPPORTED_IN_PROCESS_PYTHON_UDF"):
self.spark.range(1).select(column).collect()
conf = self.spark.sparkContext._jvm.org.apache.spark.SparkEnv.get().conf()
key = "spark.executor.pyspark.memory"
try:
conf.set(key, "128m")
# Columns are session-agnostic; validate using the query's session at execution.
identity("id")
with self.assertRaisesRegex(Exception, "UNSUPPORTED_IN_PROCESS_PYTHON_UDF"):
self.spark.range(1).select(column).collect()
finally:
conf.remove(key)
def test_configuration_uses_the_query_session(self):
from pyspark.inprocess import inprocess_udf
identity = inprocess_udf("long")(lambda x: x)
second = self.spark.newSession()
with self.sql_conf({"spark.sql.pyspark.udf.profiler": "perf"}):
column = identity("id")
self.assertEqual([r[0] for r in second.range(2).select(column).collect()], [0, 1])
with self.assertRaisesRegex(Exception, "UNSUPPORTED_IN_PROCESS_PYTHON_UDF"):
self.spark.range(2).select(column).collect()
conf = self.spark.sparkContext._jvm.org.apache.spark.SparkEnv.get().conf()
try:
conf.set("spark.executor.pyspark.memory", "0")
self.assertEqual(self.spark.range(1).select(identity("id")).first()[0], 0)
finally:
conf.remove("spark.executor.pyspark.memory")
def test_python_packages_are_not_shadowed_by_java_imports(self):
from pyspark.inprocess import inprocess_udf
def probe(x):
import sys
import pyarrow as pa
import redis
import spire
good = spire.PYTHON_PACKAGE and redis.PYTHON_PACKAGE
good = good and sys.stdout.line_buffering and sys.stdout.write_through
return pa.array([good] * len(x))
self.assertTrue(
self.spark.range(1).select(inprocess_udf("boolean")(probe)("id")).first()[0]
)
def test_result_schema_adapts_to_session_representation(self):
from pyspark.inprocess import inprocess_udf
def timestamp(x):
import pyarrow as pa
return pa.array([0] * len(x), type=pa.timestamp("us", tz="UTC"))
def nested(x):
import pyarrow as pa
# Declare the field order: newer PyArrow versions sort inferred struct fields.
return pa.array(
[{"s": ["hello"], "b": b"data"}] * len(x),
pa.struct([("s", pa.list_(pa.string())), ("b", pa.binary())]),
)
for zone in ["UTC", "Etc/UTC", "America/Los_Angeles"]:
with self.sql_conf({"spark.sql.session.timeZone": zone}):
result = self.spark.range(1).select(inprocess_udf("timestamp")(timestamp)("id"))
self.assertEqual(
result.toDF("value").selectExpr("unix_micros(value)").first()[0], 0
)
with self.sql_conf({"spark.sql.execution.arrow.useLargeVarTypes": "true"}):
result = (
self.spark.range(1)
.select(inprocess_udf("struct<s:array<string>,b:binary>")(nested)("id"))
.first()[0]
)
self.assertEqual(result.s, ["hello"])
self.assertEqual(result.b, b"data")
def test_kwargs_only_function(self):
from pyspark.inprocess import inprocess_udf
identity = inprocess_udf("long")(lambda **cols: cols["x"])
self.assertEqual(
[row[0] for row in self.spark.range(2).select(identity(x="id")).collect()], [0, 1]
)
def test_isolated_interpreter_and_explicit_process_pythonpath(self):
from pyspark.inprocess import inprocess_udf
def flags(x):
import faulthandler
import sys
import _inprocess_process_helper
import pyarrow as pa
value = (
sys.flags.isolated == 1
and sys.flags.ignore_environment == 1
and not faulthandler.is_enabled()
and _inprocess_process_helper.MAGIC == -1
)
return pa.array([value] * len(x))
self.assertTrue(
self.spark.range(1).select(inprocess_udf("boolean")(flags)("id")).first()[0]
)
@staticmethod
def _java_command(jvm, java_home):
"""
The start of a command line for a fresh probe JVM: this test's java and the module
options Spark's own launcher passes (JavaModuleOptions), so a probe gets the same module
options as Spark's JVMs on every Java version, e.g. with Unsafe allowed for Netty on
Java 25+. Only the module options are shared: the other options python/run-tests.py
gives the test JVM, such as java.io.tmpdir, are not.
The options start with -XX:+IgnoreUnrecognizedVMOptions, which also applies to the
options a probe appends, so a misspelled -XX option does not fail the launch.
"""
options = jvm.org.apache.spark.launcher.JavaModuleOptions.defaultModuleOptions()
return [str(Path(java_home) / "bin" / "java")] + options.split(" ")
def test_missing_cdi_dependency_fails_plugin_startup(self):
import subprocess
jvm = self.spark.sparkContext._jvm
java_home = jvm.java.lang.System.getProperty("java.home")
# Spark bundles arrow-c-data too, so drop its JAR as well as the configured one.
classpath = os.pathsep.join(
p
for p in jvm.java.lang.System.getProperty("java.class.path").split(os.pathsep)
if p != self.cdi_jar and not Path(p).name.startswith("arrow-c-data")
)
source = """
import java.lang.reflect.Proxy;
import java.util.Collections;
import org.apache.spark.SparkConf;
import org.apache.spark.api.plugin.PluginContext;
import org.apache.spark.sql.execution.python.InProcessPythonPlugin;
class MissingCdiProbe {
public static void main(String[] args) {
PluginContext context = (PluginContext) Proxy.newProxyInstance(
PluginContext.class.getClassLoader(), new Class<?>[] {PluginContext.class},
(proxy, method, values) -> new SparkConf(false));
try {
new InProcessPythonPlugin().executorPlugin().init(context, Collections.emptyMap());
throw new AssertionError("Plugin accepted a missing CDI dependency");
} catch (IllegalStateException expected) {
if (!(expected.getCause() instanceof NoClassDefFoundError)) throw expected;
if (!expected.getMessage().contains("arrow-c-data.jar")) throw expected;
System.out.println("MISSING_CDI_REJECTED");
}
}
}
"""
with tempfile.TemporaryDirectory() as directory:
source_file = Path(directory) / "MissingCdiProbe.java"
source_file.write_text(source)
result = subprocess.run(
self._java_command(jvm, java_home) + ["-cp", classpath, str(source_file)],
capture_output=True,
text=True,
timeout=90,
)
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
self.assertIn("MISSING_CDI_REJECTED", result.stdout)
def test_missing_jep_has_an_initialization_hint(self):
import subprocess
jvm = self.spark.sparkContext._jvm
java_home = jvm.java.lang.System.getProperty("java.home")
classpath = os.pathsep.join(
p
for p in jvm.java.lang.System.getProperty("java.class.path").split(os.pathsep)
if p != str(self.jep_jar)
)
source = """
import java.lang.reflect.Proxy;
import java.util.Collections;
import org.apache.spark.SparkConf;
import org.apache.spark.api.plugin.PluginContext;
import org.apache.spark.sql.execution.python.InProcessPythonPlugin;
import org.apache.spark.sql.execution.python.InProcessPythonRuntime;
class MissingJepProbe {
public static void main(String[] args) {
try {
InProcessPythonRuntime.currentSession();
throw new AssertionError("Uninitialized runtime was accepted");
} catch (IllegalStateException expected) {
if (!expected.getMessage().contains("executor plugin")) throw expected;
}
PluginContext context = (PluginContext) Proxy.newProxyInstance(
PluginContext.class.getClassLoader(), new Class<?>[] {PluginContext.class},
(proxy, method, values) -> new SparkConf(false));
try {
new InProcessPythonPlugin().executorPlugin().init(context, Collections.emptyMap());
throw new AssertionError("Plugin accepted a missing JEP dependency");
} catch (IllegalStateException expected) {
if (!(expected.getCause() instanceof LinkageError)) throw expected;
if (!expected.getMessage().contains("jep.jar")) throw expected;
System.out.println("MISSING_JEP_REJECTED");
}
}
}
"""
with tempfile.TemporaryDirectory() as directory:
source_file = Path(directory) / "MissingJepProbe.java"
source_file.write_text(source)
result = subprocess.run(
self._java_command(jvm, java_home)
+ [
"-cp",
classpath,
str(source_file),
],
capture_output=True,
text=True,
timeout=90,
)
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
self.assertIn("MISSING_JEP_REJECTED", result.stdout)
def test_fresh_jvm_retries_configuration_without_spark_home(self):
import subprocess
import sysconfig
import venv
from pyspark import cloudpickle
jvm = self.spark.sparkContext._jvm
java_home = jvm.java.lang.System.getProperty("java.home")
classpath = jvm.java.lang.System.getProperty("java.class.path")
spark_home = jvm.java.lang.System.getenv("SPARK_HOME")
self.assertIsNotNone(spark_home)
python_lib = Path(spark_home) / "python" / "lib"
env = os.environ.copy()
env.pop("SPARK_HOME", None)
env.pop("VIRTUAL_ENV", None)
env["PATH"] = os.defpath
# These flags must be ignored by the embedded interpreter. No signals are raised.
env["PYTHONFAULTHANDLER"] = "1"
env["PYTHONDEVMODE"] = "1"
env["LC_ALL"] = "C"
env["LANG"] = "C"
# The failed configuration also lists a default sys.path entry, which a retry keeps.
stdlib = os.path.realpath(sysconfig.get_paths()["stdlib"])
class BootstrapCheck:
def __reduce__(self):
return eval, (
"(__import__('sys').flags.isolated == 1 and "
f"{stdlib!r} in [__import__('os').path.realpath(p) "
"for p in __import__('sys').path] and "
"__import__('sys').flags.ignore_environment == 1 and "
"not __import__('faulthandler').is_enabled() and "
"'pyspark.zip' in __import__('pyspark').__file__ and "
"__import__('sys').stdout.line_buffering and "
"__import__('sys').stdout.write_through and "
"__import__('importlib.util', fromlist=['find_spec'])"
".find_spec('only_in_failed_config') is None and "
"(print('SHUTDOWN_FLUSH', end='') or (lambda x: x))) or "
"(_ for _ in ()).throw(AssertionError('unexpected bootstrap state'))",
)
source = """
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.spark.sql.execution.python.InProcessPythonRuntime;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.util.ArrowUtils;
class BootstrapProbe {
public static void main(String[] args) throws Exception {
var bad = scala.jdk.javaapi.CollectionConverters
.asScala(Arrays.asList(args[0].split(java.io.File.pathSeparator))).toSeq();
boolean failed = false;
try {
InProcessPythonRuntime.initialize(bad);
} catch (jep.JepException expected) {
failed = true;
}
if (!failed) throw new AssertionError("Expected missing JEP package");
// A configured directory that did not exist when the first attempt failed.
Files.createSymbolicLink(Path.of(args[5]), Path.of(args[2]));
var good = scala.jdk.javaapi.CollectionConverters
.asScala(Arrays.asList(args[1], args[5])).toSeq();
try {
InProcessPythonRuntime.initialize(good);
Field field = ArrowUtils.toArrowField("result", DataTypes.LongType, true, "UTC",
false, org.apache.spark.sql.types.Metadata.empty(), false);
InProcessPythonRuntime.currentSession().register("probe",
Files.readAllBytes(Path.of(args[3])), field, args[4], false, false, false, true);
if (ArrowUtils.rootAllocator().getAllocatedMemory() != 0) {
throw new AssertionError("Unreleased registration schema");
}
System.out.println("BOOTSTRAP_OK");
} finally {
InProcessPythonRuntime.currentSession().release(
scala.jdk.javaapi.CollectionConverters.asScala(Arrays.asList("probe")).toSeq());
InProcessPythonRuntime.shutdown();
}
}
}
"""
with tempfile.TemporaryDirectory() as directory:
# Keep JEP absent initially even when it is installed in the system Python.
clean_python = Path(directory) / "python"
venv.EnvBuilder(with_pip=False).create(clean_python)
env["PATH"] = str(clean_python / "bin") + os.pathsep + os.defpath
# A module that only the failed configuration's directory has, which the retried
# configuration must not import.
(Path(directory) / "only_in_failed_config.py").write_text("")
late = Path(directory) / "created_after_failure"
# Without a packaged zip, e.g. in CI, package the sources for this process.
packaged = python_lib / "pyspark.zip"
if not packaged.is_file():
packaged = Path(directory) / "pyspark.zip"
with zipfile.ZipFile(packaged, "w") as archive:
sources = Path(spark_home) / "python" / "pyspark"
for entry in sources.rglob("*"):
if entry.is_file() and "__pycache__" not in entry.parts:
archive.write(entry, entry.relative_to(sources.parent))
env["PYTHONPATH"] = os.pathsep.join(
[str(packaged)] + [str(p) for p in python_lib.glob("py4j-*.zip")]
)
source_file = Path(directory) / "BootstrapProbe.java"
source_file.write_text(source)
command_file = Path(directory) / "command.pickle"
command_file.write_bytes(cloudpickle.dumps(BootstrapCheck()))
import sys
result = subprocess.run(
self._java_command(jvm, java_home)
+ [
f"-Djava.library.path={self.jep_dir}",
"--class-path",
classpath,
str(source_file),
os.pathsep.join([directory, stdlib, str(late)]),
self.site_packages,
str(self.jep_dir.parent),
str(command_file),
"%d.%d" % sys.version_info[:2],
str(late),
],
env=env,
capture_output=True,
text=True,
timeout=90,
)
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
self.assertIn("BOOTSTRAP_OK", result.stdout)
self.assertIn("SHUTDOWN_FLUSH", result.stdout)
self.assertIn("set LC_ALL=C.UTF-8", result.stderr)
def test_failed_bootstrap_preserves_message_and_freezes_site_packages(self):
import subprocess
jvm = self.spark.sparkContext._jvm
java_home = jvm.java.lang.System.getProperty("java.home")
classpath = jvm.java.lang.System.getProperty("java.class.path")
source = """
import java.util.Arrays;
import org.apache.spark.sql.execution.python.InProcessPythonRuntime;
class BootstrapFailureProbe {
public static void main(String[] args) {
var paths = scala.jdk.javaapi.CollectionConverters
.asScala(Arrays.asList(args[0], args[1])).toSeq();
try {
InProcessPythonRuntime.initialize(paths);
throw new AssertionError("Expected an unsupported PyArrow version");
} catch (jep.JepException expected) {
String message = expected.getMessage();
if (!message.contains("PySparkImportError") ||
!message.contains("UNSUPPORTED_PACKAGE_VERSION") ||
!message.contains("PyArrow") || !message.contains("0.0.0")) throw expected;
}
var changed = scala.jdk.javaapi.CollectionConverters
.asScala(Arrays.asList(args[1])).toSeq();
try {
InProcessPythonRuntime.initialize(changed);
throw new AssertionError("Accepted a different environment after bootstrap failure");
} catch (IllegalStateException expected) {
if (!expected.getMessage().contains("Restart the executor process")) throw expected;
}
System.out.println("BOOTSTRAP_FAILURE_CHECKED");
}
}
"""
with tempfile.TemporaryDirectory() as directory:
# Use a backslash in a supported POSIX path to exercise JEP's escaping too.
packages = Path(directory) / "back\\slash"
packages.mkdir()
(packages / "old_arrow.pth").write_text(
"import pyarrow; pyarrow.__version__ = '0.0.0'\n"
)
source_file = Path(directory) / "BootstrapFailureProbe.java"
source_file.write_text(source)
env = os.environ.copy()
# This direct JVM probe needs no Spark installation and always imports source.
# Py4J comes from the same place as in this process, e.g. Spark's source zip.
import py4j
env["SPARK_HOME"] = directory
env["PYTHONPATH"] = os.pathsep.join(
[str(self.python_source), str(Path(py4j.__file__).parents[1])]
)
result = subprocess.run(
self._java_command(jvm, java_home)
+ [
f"-Djava.library.path={self.jep_dir}",
"--class-path",
classpath,
str(source_file),
str(packages),
str(self.jep_dir.parent),
],
env=env,
capture_output=True,
text=True,
timeout=90,
)
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
self.assertIn("BOOTSTRAP_FAILURE_CHECKED", result.stdout)
def test_exception_unicode_and_nul_survive_jep(self):
from pyspark.inprocess import inprocess_udf
def fail(x):
raise ValueError("failure: caf\u00e9 \u4e2d\u6587 \U0001f600 \ud800 \0 tail")
with self.assertRaises(Exception) as error:
self.spark.range(1).select(inprocess_udf("long")(fail)("id")).collect()
self.assertIn("caf\u00e9 \u4e2d\u6587", str(error.exception))
self.assertIn(r"\U0001f600 \ud800 \x00 tail", str(error.exception))
def test_named_argument_resolver(self):
from pyspark.inprocess import inprocess_udf
identity = inprocess_udf("long")(lambda x, **kw: x)
with self.sql_conf({"spark.sql.caseSensitive": "false"}):
with self.assertRaisesRegex(Exception, "DOUBLE_NAMED_ARGUMENT_REFERENCE"):
identity(x="id", X="id")
with self.sql_conf({"spark.sql.caseSensitive": "true"}):
result = self.spark.range(2).select(identity(x="id", X="id"))
self.assertEqual([r[0] for r in result.collect()], [0, 1])
def test_spark_python_distribution_precedes_site_packages(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
def location(x):
import pyspark
return pa.array([pyspark.__file__] * len(x))
path = self.spark.range(1).select(inprocess_udf("string")(location)("id")).first()[0]
expected = str(self.python_source / "pyspark" / "__init__.py")
packaged = str(self.python_source / "lib" / "pyspark.zip" / "pyspark" / "__init__.py")
self.assertIn(path, [expected, packaged])
def test_declared_struct_metadata(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType, StructField, StructType
declared = StructType([StructField("x", LongType(), metadata={"comment": "c"})])
identity = inprocess_udf(declared)(lambda x: x)
df = self.spark.sql("SELECT named_struct('x', 7L) AS value")
result = df.select(identity("value").alias("value"))
self.assertEqual(result.collect(), df.collect())
self.assertEqual(result.schema[0].dataType, declared)
def test_temporal_precision_metadata(self):
from pyspark.inprocess import inprocess_udf
for declared in ["time(3)", "timestamp_ntz(7)", "timestamp_ltz(8)"]:
literal = "12:34:56.123" if declared.startswith("time(") else "2024-01-02 12:34:56.123"
for nested in [False, True]:
with self.subTest(declared=declared, nested=nested):
df = self.spark.sql(f"SELECT CAST('{literal}' AS {declared}) AS value")
if nested:
df = df.selectExpr("named_struct('t', value) AS value")
identity = inprocess_udf(df.schema[0].dataType)(lambda x: x)
result = df.select(identity("value").alias("value"))
self.assertEqual(result.schema[0].dataType, df.schema[0].dataType)
self.assertEqual(
result.selectExpr("CAST(value AS STRING)").collect(),
df.selectExpr("CAST(value AS STRING)").collect(),
)
def test_large_types_variant_and_spatial_identity(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import Geography, GeographyType, Geometry, GeometryType
wkb = bytes.fromhex("010100000000000000000031400000000000001c40")
frames = [
self.spark.sql("SELECT parse_json('{\"a\":1}') AS value"),
self.spark.createDataFrame([(Geometry(wkb, 0),)], "value geometry(0)"),
self.spark.createDataFrame([(Geography(wkb, 4326),)], "value geography(4326)"),
]
for large in ["false", "true"]:
with self.sql_conf({"spark.sql.execution.arrow.useLargeVarTypes": large}):
for df in frames:
with self.subTest(large=large, data_type=df.schema[0].dataType):
identity = inprocess_udf(df.schema[0].dataType)(lambda x: x)
result = df.select(identity("value").alias("value"))
if isinstance(df.schema[0].dataType, (GeometryType, GeographyType)):
self.assertEqual(result.collect(), df.collect())
else:
self.assertEqual(
result.selectExpr("CAST(value AS STRING)").collect(),
df.selectExpr("CAST(value AS STRING)").collect(),
)
def test_ddl_return_type_and_nondeterminism(self):
from pyspark.inprocess import inprocess_udf
identity = inprocess_udf("long")(lambda x: x).asNondeterministic()
self.assertEqual(self.spark.range(2).select(identity("id")).first()[0], 0)
plan = self.spark.range(1).select(identity("id"))._jdf.queryExecution().analyzed()
self.assertFalse(plan.expressions().apply(0).deterministic())
def test_command_is_stable_after_the_first_call(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
# Like other Python UDFs, the command is serialized once, on the first call.
namespace = {"inprocess_udf": inprocess_udf, "pa": pa}
exec(
"@inprocess_udf('long')\n"
"def f(x): return pa.array([LOOKUP] * len(x), type=pa.int64())\n",
namespace,
)
wrapper = namespace["f"]
namespace["LOOKUP"] = 42
first = wrapper("id")
namespace["LOOKUP"] = 99
rows = self.spark.range(1).select(first, wrapper("id")).collect()
self.assertEqual(tuple(rows[0]), (42, 42))
def test_calls_share_one_jvm_function(self):
import copy
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
jvm = self.spark._jvm
double = inprocess_udf("long")(lambda x: pc.multiply(x, 2))
def functions(*cols):
plan = self.spark.range(1).select(*cols)._jdf.queryExecution().analyzed()
udfs = [plan.expressions().apply(i).child() for i in range(len(cols))]
return [jvm.System.identityHashCode(udf.func()) for udf in udfs], udfs
# The command crosses py4j once, however many expressions call the UDF.
ids, _ = functions(double("id"), double("id"))
self.assertEqual(ids[0], ids[1])
# Nondeterminism belongs to each call, so it keeps the function.
double.asNondeterministic()
later, udfs = functions(double("id"))
self.assertEqual(later[0], ids[0])
self.assertFalse(udfs[0].udfDeterministic())
# A called wrapper can still be copied; the copy creates its own function.
copied = copy.deepcopy(double)
self.assertIsNone(copied._jfunc)
self.assertEqual(self.spark.range(2).select(copied("id")).collect()[1][0], 2)
def test_traceback_settings_are_per_registration(self):
from pyspark.inprocess import inprocess_udf
# User frames must be outside the pyspark package for the worker's simplifier.
namespace = {}
exec(
"def probe(x):\n probe_local = 8675309\n"
" raise ValueError('traceback policy probe')",
namespace,
)
traceback_probe = inprocess_udf("long")(namespace["probe"])
for hide, simplified, locals_enabled in [
(True, False, False),
(False, True, False),
(False, False, False),
(True, True, True),
(False, True, True),
(False, False, True),
]:
with self.sql_conf(
{
"spark.sql.execution.pyspark.udf.tracebackWithLocals.enabled": str(
locals_enabled
).lower(),
"spark.sql.execution.pyspark.udf.hideTraceback.enabled": str(hide).lower(),
"spark.sql.execution.pyspark.udf.simplifiedTraceback.enabled": str(
simplified
).lower(),
}
):
with self.assertRaises(Exception) as error:
self.spark.range(1).select(traceback_probe("id")).collect()
message = str(error.exception)
self.assertIn("ValueError: traceback policy probe", message)
self.assertEqual('File "' in message, not hide)
self.assertEqual("probe_local = 8675309" in message, locals_enabled and not hide)
if not hide:
self.assertEqual("inprocess/runtime.py" in message, not simplified)
def test_embedded_hash_seed_matches_default_worker_seed(self):
import subprocess
import sys
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
expected = int(
subprocess.check_output(
[sys.executable, "-c", "print(hash('spark'))"],
env={**os.environ, "PYTHONHASHSEED": "0"},
)
)
hash_udf = inprocess_udf("long")(
lambda x: pa.array([hash("spark")] * len(x), type=pa.int64())
)
rows = self.spark.range(4, numPartitions=2).select(hash_udf("id")).collect()
self.assertEqual([row[0] for row in rows], [expected] * 4)
def test_expression_arguments_and_multiple_batches(self):
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
from pyspark.sql.functions import lit
from pyspark.sql.types import LongType
add = inprocess_udf(LongType())(lambda x, y: pc.add(x, y))
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "2"}):
df = self.spark.range(11, numPartitions=3)
values = df.select(add(df.id + 1, lit(2).cast("long"))).collect()
self.assertEqual([r[0] for r in values], list(range(3, 14)))
def test_preserved_child_columns_produce_collectable_rows(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType())(lambda x: x)
df = self.spark.range(3)
rows = df.select(df.id, identity(df.id)).collect()
self.assertEqual([tuple(r) for r in rows], [(0, 0), (1, 1), (2, 2)])
def test_unlimited_batch_size(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType())(lambda x: x)
for batch_size in (0, -1):
with (
self.subTest(batch_size=batch_size),
self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": str(batch_size)}),
):
df = self.spark.range(5)
values = [r[0] for r in df.select(identity(df.id)).collect()]
self.assertEqual(values, list(range(5)))
def test_byte_limit_applies_without_a_row_limit(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
lengths = inprocess_udf(LongType())(lambda x: pa.array([len(x)] * len(x), type=pa.int64()))
with self.sql_conf(
{
"spark.sql.execution.arrow.maxRecordsPerBatch": "0",
"spark.sql.execution.arrow.maxBytesPerBatch": "1",
}
):
df = self.spark.range(4)
self.assertEqual([r[0] for r in df.select(lengths(df.id)).collect()], [1] * 4)
def test_wrong_result_length_fails_before_rows_are_read(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
short = inprocess_udf(LongType())(lambda x: x.slice(0, len(x) - 1))
df = self.spark.range(3, numPartitions=1)
with self.assertRaisesRegex(Exception, "returned 2 rows; expected 3"):
df.select(short(df.id)).collect()
def test_numpy_finalizers_run_on_the_interpreter_thread(self):
from pyspark.inprocess import inprocess_udf
with tempfile.TemporaryDirectory() as directory:
marker = str(Path(directory) / "finalizers.txt")
def produce(x):
import threading
import weakref
import numpy as np
import pyarrow as pa
owner = threading.get_ident()
def finalized():
with open(marker, "a") as stream:
stream.write(f"{owner} {threading.get_ident()}\n")
values = np.arange(len(x), dtype=np.int64)
weakref.finalize(values, finalized)
return pa.array(values)
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "2"}):
result = (
self.spark.range(8, numPartitions=2)
.select(inprocess_udf("long")(produce)("id"))
.collect()
)
self.assertEqual(len(result), 8)
# A subsequent invocation waits behind cleanup already queued by completed tasks.
identity = inprocess_udf("long")(lambda x: x)
self.spark.range(1).select(identity("id")).collect()
records = Path(marker).read_text().splitlines()
self.assertEqual(len(records), 4)
for record in records:
owner, finalizer = record.split()
self.assertEqual(owner, finalizer)
def test_arrow_memory_is_released_on_success_limit_and_failure(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType())(lambda x: x)
@inprocess_udf(LongType())
def fail(x):
raise ValueError("second UDF failed")
arrow_utils = self.spark.sparkContext._jvm.org.apache.spark.sql.util.ArrowUtils
allocator = arrow_utils.rootAllocator()
before = allocator.getAllocatedMemory()
# The interpreter thread releases Python's references after the query returns, and a
# failed job does not wait for its other tasks to release their memory.
@eventually(timeout=30.0, catch_assertions=True)
def assert_released():
self.assertEqual(allocator.getAllocatedMemory(), before)
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "2"}):
df = self.spark.range(9, numPartitions=3)
for _ in range(3):
self.assertEqual(df.select(identity(df.id)).collect()[0][0], 0)
assert_released()
# The filter keeps the limit above the in-process node, which then stops early.
df.select(identity(df.id).alias("v")).filter("v >= 0").limit(1).collect()
assert_released()
with self.assertRaisesRegex(Exception, "second UDF failed"):
df.select(identity(df.id), fail(df.id)).collect()
assert_released()
# Deserialization fails before Python imports any input CDI structures.
broken = inprocess_udf(LongType())(lambda x: x)
with patch("pyspark.inprocess.udf._serialize_udf", return_value=b"invalid pickle"):
column = broken(df.id)
with self.assertRaisesRegex(Exception, "UnpicklingError"):
df.select(column).collect()
assert_released()
def test_double_long(self):
"""@inprocess_udf with LongType input/output doubles each value."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType
@inprocess_udf(returnType=LongType())
def double(x):
return pc.multiply(x, 2)
df = self.spark.range(1, 6) # [1, 2, 3, 4, 5]
result = df.select(double(df["id"])).collect()
self.assertEqual([r[0] for r in result], [2, 4, 6, 8, 10])
def test_negate_double(self):
"""@inprocess_udf with DoubleType negates each value."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import DoubleType
@inprocess_udf(returnType=DoubleType())
def negate(x):
return pc.negate(x)
data = [(1.5,), (2.5,), (3.0,)]
df = self.spark.createDataFrame(data, ["v"])
result = [r[0] for r in df.select(negate(df["v"])).collect()]
self.assertAlmostEqual(result[0], -1.5)
self.assertAlmostEqual(result[1], -2.5)
self.assertAlmostEqual(result[2], -3.0)
def test_identity_integer(self):
"""@inprocess_udf with IntegerType passes values through unchanged."""
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import IntegerType
@inprocess_udf(returnType=IntegerType())
def identity(x):
return x
data = [(i,) for i in range(5)]
df = self.spark.createDataFrame(data, "v int")
result = [r[0] for r in df.select(identity(df["v"])).collect()]
self.assertEqual(result, list(range(5)))
def test_boolean_not(self):
"""@inprocess_udf with BooleanType inverts each boolean."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import BooleanType
@inprocess_udf(returnType=BooleanType())
def invert(x):
return pc.invert(x)
data = [(True,), (False,), (True,)]
df = self.spark.createDataFrame(data, ["v"])
result = [r[0] for r in df.select(invert(df["v"])).collect()]
self.assertEqual(result, [False, True, False])
# ------------------------------------------------------------------
# Null handling
# ------------------------------------------------------------------
def test_null_passthrough(self):
"""Null values in the input must produce null in the output."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType
@inprocess_udf(returnType=LongType())
def negate(x):
return pc.negate(x)
data = [1, None, 3]
df = self.spark.createDataFrame([(v,) for v in data], ["v"])
rows = df.select(negate(df["v"])).collect()
self.assertEqual(rows[0][0], -1)
self.assertIsNone(rows[1][0])
self.assertEqual(rows[2][0], -3)
def test_all_nulls(self):
"""Column of all-null values: every output row must be null."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType, StructField, StructType
@inprocess_udf(returnType=LongType())
def double(x):
return pc.multiply(x, 2)
schema = StructType([StructField("v", LongType(), nullable=True)])
data = [(None,), (None,), (None,)]
df = self.spark.createDataFrame(data, schema)
rows = df.select(double(df["v"])).collect()
for row in rows:
self.assertIsNone(row[0])
# ------------------------------------------------------------------
# Multi-column UDFs
# ------------------------------------------------------------------
def test_two_column_add(self):
"""UDF that adds two LongType columns together."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType
@inprocess_udf(returnType=LongType())
def add(a, b):
return pc.add(a, b)
data = [(1, 10), (2, 20), (3, 30)]
df = self.spark.createDataFrame(data, ["a", "b"])
result = [r[0] for r in df.select(add(df["a"], df["b"])).collect()]
self.assertEqual(result, [11, 22, 33])
def test_two_column_multiply(self):
"""UDF that multiplies two DoubleType columns."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import DoubleType
@inprocess_udf(returnType=DoubleType())
def multiply(a, b):
return pc.multiply(a, b)
data = [(2.0, 3.0), (4.0, 5.0)]
df = self.spark.createDataFrame(data, ["a", "b"])
result = [r[0] for r in df.select(multiply(df["a"], df["b"])).collect()]
self.assertAlmostEqual(result[0], 6.0)
self.assertAlmostEqual(result[1], 20.0)
# ------------------------------------------------------------------
# UDF reuse and multiple UDFs on the same query
# ------------------------------------------------------------------
def test_two_udfs_same_select(self):
"""Two different @inprocess_udf calls in the same select are both executed."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType
@inprocess_udf(returnType=LongType())
def double(x):
return pc.multiply(x, 2)
@inprocess_udf(returnType=LongType())
def triple(x):
return pc.multiply(x, 3)
df = self.spark.range(1, 4) # [1, 2, 3]
rows = df.select(double(df["id"]), triple(df["id"])).collect()
self.assertEqual([r[0] for r in rows], [2, 4, 6])
self.assertEqual([r[1] for r in rows], [3, 6, 9])
def test_udf_reuse_across_queries(self):
"""The same InProcessUDFWrapper can be applied to different DataFrames."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType
@inprocess_udf(returnType=LongType())
def double(x):
return pc.multiply(x, 2)
df1 = self.spark.range(1, 4)
df2 = self.spark.range(10, 13)
result1 = [r[0] for r in df1.select(double(df1["id"])).collect()]
result2 = [r[0] for r in df2.select(double(df2["id"])).collect()]
self.assertEqual(result1, [2, 4, 6])
self.assertEqual(result2, [20, 22, 24])
# ------------------------------------------------------------------
# Concurrency
# ------------------------------------------------------------------
def test_concurrent_tasks_have_separate_function_state(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
state = []
@inprocess_udf(LongType(), deterministic=False)
def counter(x):
state.append(1)
return pa.array([len(state)] * len(x), type=pa.int64())
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "2"}):
df = self.spark.range(8, numPartitions=2)
for _ in range(2):
self.assertEqual(
[r[0] for r in df.select(counter(df.id)).collect()],
[1, 1, 2, 2, 1, 1, 2, 2],
)
# ------------------------------------------------------------------
# Closure capture
# ------------------------------------------------------------------
def test_udf_captures_closure(self):
"""UDF closure values defined in outer scope are serialized correctly."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType
factor = 7 # captured in closure
@inprocess_udf(returnType=LongType())
def scale(x):
return pc.multiply(x, factor)
df = self.spark.range(1, 4)
result = [r[0] for r in df.select(scale(df["id"])).collect()]
self.assertEqual(result, [7, 14, 21])
# ------------------------------------------------------------------
# Non-deterministic flag
# ------------------------------------------------------------------
def test_nondeterministic_udf_executes(self):
"""A UDF declared deterministic=False executes and returns correct values."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType
@inprocess_udf(returnType=LongType(), deterministic=False)
def double(x):
return pc.multiply(x, 2)
df = self.spark.range(1, 4)
result = [r[0] for r in df.select(double(df["id"])).collect()]
self.assertEqual(result, [2, 4, 6])
def test_nondeterministic_flag_propagates_to_expression(self):
"""deterministic=False must be reflected in the PythonUDF expression."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType
@inprocess_udf(returnType=LongType(), deterministic=False)
def double(x):
return pc.multiply(x, 2)
df = self.spark.range(3)
jdf = df.select(double(df["id"]))._jdf
# Python UDF extraction happens during optimization.
optimized = jdf.queryExecution().optimizedPlan()
# Walk the logical plan via children() (no PartialFunction needed)
# to locate the ArrowEvalPython node inserted during optimization.
def find_node(plan):
if plan.getClass().getSimpleName() == "ArrowEvalPython":
return plan
children = plan.children().toList()
for i in range(children.length()):
found = find_node(children.apply(i))
if found is not None:
return found
return None
inprocess_node = find_node(optimized)
self.assertIsNotNone(inprocess_node, "ArrowEvalPython not found in analyzed plan")
udfs = inprocess_node.udfs().toList()
self.assertGreater(udfs.length(), 0)
self.assertFalse(
udfs.apply(0).deterministic(),
"PythonUDF with deterministic=False must have deterministic()==False",
)
# ------------------------------------------------------------------
# String type
# ------------------------------------------------------------------
def test_string_identity(self):
"""@inprocess_udf with StringType passes strings through unchanged."""
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import StringType
@inprocess_udf(returnType=StringType())
def identity(s):
return s
data = [("hello",), ("world",), (None,)]
df = self.spark.createDataFrame(data, ["v"])
result = [r[0] for r in df.select(identity(df["v"])).collect()]
self.assertEqual(result[0], "hello")
self.assertEqual(result[1], "world")
self.assertIsNone(result[2])
def test_string_upper(self):
"""@inprocess_udf with StringType applies utf8_upper transformation."""
import pyarrow.compute as pc
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import StringType
@inprocess_udf(returnType=StringType())
def upper(s):
return pc.utf8_upper(s)
data = [("hello",), ("world",)]
df = self.spark.createDataFrame(data, ["v"])
result = [r[0] for r in df.select(upper(df["v"])).collect()]
self.assertEqual(result, ["HELLO", "WORLD"])
# ------------------------------------------------------------------
# Binary type
# ------------------------------------------------------------------
def test_binary_identity(self):
"""@inprocess_udf with BinaryType passes bytes through unchanged."""
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import BinaryType
@inprocess_udf(returnType=BinaryType())
def identity(b):
return b
data = [(b"hello",), (b"world",), (None,)]
df = self.spark.createDataFrame(data, ["v"])
result = [r[0] for r in df.select(identity(df["v"])).collect()]
self.assertEqual(bytes(result[0]), b"hello")
self.assertEqual(bytes(result[1]), b"world")
self.assertIsNone(result[2])
# ------------------------------------------------------------------
# Array type
# ------------------------------------------------------------------
def test_array_identity(self):
"""@inprocess_udf with ArrayType(LongType()) passes arrays through unchanged."""
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import ArrayType, LongType
@inprocess_udf(returnType=ArrayType(LongType()))
def identity(arr):
return arr
data = [([1, 2, 3],), ([4, 5],), (None,)]
df = self.spark.createDataFrame(data, ["v"])
result = [r[0] for r in df.select(identity(df["v"])).collect()]
self.assertEqual(list(result[0]), [1, 2, 3])
self.assertEqual(list(result[1]), [4, 5])
self.assertIsNone(result[2])
# ------------------------------------------------------------------
# Struct type
# ------------------------------------------------------------------
def test_struct_identity(self):
"""@inprocess_udf with StructType passes structs through unchanged."""
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import DoubleType, LongType, StructField, StructType
inner = StructType([StructField("a", LongType()), StructField("b", DoubleType())])
outer = StructType([StructField("v", inner)])
@inprocess_udf(returnType=inner)
def identity(s):
return s
data = [((1, 2.0),), ((3, 4.0),)]
df = self.spark.createDataFrame(data, outer)
result = [r[0] for r in df.select(identity(df["v"])).collect()]
self.assertEqual(result[0]["a"], 1)
self.assertAlmostEqual(result[0]["b"], 2.0)
self.assertEqual(result[1]["a"], 3)
self.assertAlmostEqual(result[1]["b"], 4.0)
# ------------------------------------------------------------------
# Date type
# ------------------------------------------------------------------
def test_date_identity(self):
"""@inprocess_udf with DateType passes dates through unchanged."""
import datetime
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import DateType
@inprocess_udf(returnType=DateType())
def identity(d):
return d
dates = [datetime.date(2024, 1, 1), datetime.date(2024, 6, 15), None]
data = [(d,) for d in dates]
df = self.spark.createDataFrame(data, ["v"])
result = [r[0] for r in df.select(identity(df["v"])).collect()]
self.assertEqual(result[0], datetime.date(2024, 1, 1))
self.assertEqual(result[1], datetime.date(2024, 6, 15))
self.assertIsNone(result[2])
# ------------------------------------------------------------------
# Timestamp type
# ------------------------------------------------------------------
def test_timestamp_identity(self):
"""@inprocess_udf with TimestampType passes timestamps through unchanged."""
import datetime
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import TimestampType
@inprocess_udf(returnType=TimestampType())
def identity(ts):
return ts
data = [(datetime.datetime(2024, 3, 15, 10, 30, 0),), (None,)]
df = self.spark.createDataFrame(data, ["v"])
result = [r[0] for r in df.select(identity(df["v"])).collect()]
self.assertIsNone(result[1])
# Check date components are preserved (timezone handling may shift hours)
self.assertEqual(result[0].year, 2024)
self.assertEqual(result[0].month, 3)
self.assertEqual(result[0].day, 15)
# ------------------------------------------------------------------
# sitePackages / sys.path extension
# ------------------------------------------------------------------
def test_site_packages_path_extension_works_in_interpreter(self):
"""The plugin's actual sitePackages config exposes a module to the interpreter."""
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
@inprocess_udf(LongType())
def read_magic(x):
import _inprocess_test_helper
import pyarrow as pa
return pa.array([_inprocess_test_helper.MAGIC] * len(x), type=pa.int64())
df = self.spark.range(1)
self.assertEqual(df.select(read_magic(df.id)).first()[0], 99)
def test_bootstrap_converts_base_exceptions_and_retries_the_same_configuration(self):
jvm = self.spark.sparkContext._jvm
runtime = jvm.org.apache.spark.sql.execution.python.InProcessPythonRuntime
runtime_module = getattr(
getattr(jvm.org.apache.spark.sql.execution.python, "InProcessPythonRuntime$"), "MODULE$"
)
def initialize():
paths = jvm.java.util.ArrayList()
paths.add(self.site_packages)
paths.add(str(self.jep_dir.parent))
runtime.initialize(jvm.PythonUtils.toSeq(paths))
# Check the shared guard independently of CPython/JEP exception handling.
for script in [
"raise KeyboardInterrupt('bootstrap probe')",
"import missing_bootstrap_probe",
]:
with self.assertRaisesRegex(RuntimeError, "bootstrap failed"):
exec(runtime_module.bootstrapScript(script), {})
runtime.shutdown()
probe = Path(self.site_packages) / "probe.pth"
try:
probe.write_text("import sys; raise KeyboardInterrupt('bootstrap probe')\n")
with self.assertRaisesRegex(Exception, "bootstrap failed.*bootstrap probe"):
initialize()
finally:
probe.unlink()
initialize()
from pyspark.inprocess import inprocess_udf
identity = inprocess_udf("long")(lambda x: x)
self.assertEqual(self.spark.range(1).select(identity("id")).first()[0], 0)
def test_runtime_restart_from_driver_thread(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType())(lambda x: x)
df = self.spark.range(3)
expected = df.select(identity(df.id)).collect()
runtime = self.spark.sparkContext._jvm.org.apache.spark.sql.execution.python
runtime.InProcessPythonRuntime.shutdown()
try:
with self.assertRaisesRegex(Exception, "has been stopped"):
df.select(identity(df.id)).collect()
changed = self.spark.sparkContext._jvm.java.util.ArrayList()
changed.add(self.site_packages)
with self.assertRaisesRegex(Exception, "Restart the executor process"):
runtime.InProcessPythonRuntime.initialize(
self.spark.sparkContext._jvm.PythonUtils.toSeq(changed)
)
finally:
paths = self.spark.sparkContext._jvm.java.util.ArrayList()
paths.add(self.site_packages)
paths.add(str(self.jep_dir.parent))
runtime.InProcessPythonRuntime.initialize(
self.spark.sparkContext._jvm.PythonUtils.toSeq(paths)
)
self.assertEqual(df.select(identity(df.id)).collect(), expected)
# ------------------------------------------------------------------
# Error handling
# ------------------------------------------------------------------
def test_buggy_udf_exposes_python_traceback(self):
"""A UDF that raises an exception must include the Python traceback in the error.
The traceback must name the exception type, the error message, and the
file/line where the exception was raised, matching what you would see in a
standard Python traceback.
"""
from pyspark.inprocess.udf import inprocess_udf
from pyspark.sql.types import LongType
@inprocess_udf(returnType=LongType())
def always_fails(x):
raise ValueError("intentional test error from always_fails")
df = self.spark.range(1)
try:
df.select(always_fails(df["id"])).collect()
self.fail("Expected exception was not raised")
except Exception as e:
error_msg = str(e)
self.assertIn(
"ValueError", error_msg, "Exception type must appear in the error message"
)
self.assertIn(
"intentional test error from always_fails",
error_msg,
"Exception message must appear in the error",
)
self.assertIn(
"always_fails", error_msg, "UDF function name must appear in the traceback"
)
def test_sliced_results_are_read_correctly_by_arrow_java(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import (
ArrayType,
BooleanType,
LongType,
StringType,
StructField,
StructType,
)
cases = [
(pa.array([9, 1, None, 3]).slice(1), LongType()),
(pa.array([False, True, None, False]).slice(1), BooleanType()),
(pa.array(["discard", "one", None, "three"]).slice(1), StringType()),
(pa.array([[9], [1], None, [3]]).slice(1), ArrayType(LongType())),
(
pa.StructArray.from_arrays([pa.array([9, 1, None, 3]).slice(1)], names=["x"]),
StructType([StructField("x", LongType())]),
),
]
df = self.spark.range(3, numPartitions=1)
for value, return_type in cases:
with self.subTest(return_type=return_type):
identity = inprocess_udf(return_type)(lambda x: value)
actual = [r[0] for r in df.select(identity(df.id)).collect()]
if isinstance(return_type, StructType):
actual = [r.asDict() for r in actual]
self.assertEqual(actual, value.to_pylist())
def test_retained_inputs_are_not_overwritten_by_later_batches(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
retained = []
@inprocess_udf(LongType(), deterministic=False)
def remember(x):
for previous, snapshot in retained:
if previous.to_pylist() != snapshot:
raise ValueError("retained input changed")
retained.append((x, x.to_pylist()))
return x
df = self.spark.createDataFrame([(1,), (None,), (3,), (4,), (None,), (6,)], "x long")
df = df.coalesce(1)
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "2"}):
self.assertEqual(
[r[0] for r in df.select(remember("x")).collect()], [1, None, 3, 4, None, 6]
)
def test_pass_through_struct_with_duplicate_names(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql import functions as F
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType())(lambda x: x)
df = self.spark.range(3).select(
"id", F.struct(F.col("id"), F.col("id").cast("string").alias("id")).alias("s")
)
# Materialize the struct so CollapseProject cannot move it above the UDF.
df.cache()
try:
rows = df.select("s", identity("id")).collect()
self.assertEqual(
[(tuple(r[0]), r[1]) for r in rows], [((i, str(i)), i) for i in range(3)]
)
finally:
df.unpersist()
def test_pass_through_interval_does_not_require_arrow_conversion(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType())(lambda x: x)
df = self.spark.range(3).selectExpr(
"id", "make_interval(0, 0, 0, 0, 0, 0, 10000000000 + id) AS c"
)
df.cache()
try:
# CalendarInterval cannot be converted to a Python Row. Compare JVM Row values.
expected = df.select("c", "id")._jdf.collect()
actual = df.select("c", identity("id"))._jdf.collect()
self.assertEqual(
[(r.get(0).toString(), r.getLong(1)) for r in actual],
[(r.get(0).toString(), r.getLong(1)) for r in expected],
)
finally:
df.unpersist()
def test_mixed_worker_and_inprocess_udfs(self):
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
from pyspark.sql.functions import udf
from pyspark.sql.types import LongType
double = inprocess_udf(LongType())(lambda x: pc.multiply(x, 2))
plus_one = udf(lambda x: x + 1, LongType(), useArrow=False)
df = self.spark.range(4)
rows = df.select(double(plus_one("id")), plus_one(double("id"))).collect()
self.assertEqual([tuple(r) for r in rows], [(2 * (i + 1), 2 * i + 1) for i in range(4)])
def test_pipelined_arrow_worker_consumes_inprocess_results(self):
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
from pyspark.sql.functions import arrow_udf
conf = self.spark.sparkContext._jvm.org.apache.spark.SparkEnv.get().conf()
key = "spark.python.udf.pipelined.enabled"
previous = conf.get(key, "false")
conf.set(key, "true")
try:
double = inprocess_udf("long")(lambda x: pc.multiply(x, 2))
triple = inprocess_udf("long")(lambda x: pc.multiply(x, 3))
add = arrow_udf(lambda x, y: pc.add(x, y), "long")
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "2"}):
query = self.spark.range(8, numPartitions=2).select(add(double("id"), triple("id")))
plan = query._jdf.queryExecution().executedPlan().toString()
self.assertIn("InProcessArrowEvalPython", plan)
self.assertIn("ArrowEvalPython", plan)
self.assertEqual([r[0] for r in query.collect()], [i * 5 for i in range(8)])
finally:
conf.set(key, previous)
def test_pipelined_worker_reads_materialized_nested_results_across_batches(self):
import pyarrow as pa
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
from pyspark.sql.functions import arrow_udf
conf = self.spark.sparkContext._jvm.org.apache.spark.SparkEnv.get().conf()
key = "spark.python.udf.pipelined.enabled"
previous = conf.get(key, "false")
conf.set(key, "true")
try:
# The worker's writer thread consumes these rows; each must outlive its batch.
pair = inprocess_udf("array<string>")(
lambda s: pa.array([[v, v] for v in s.to_pylist()], pa.list_(pa.string()))
)
joined = arrow_udf(lambda a: pc.binary_join(a, "-"), "string")
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "3"}):
df = self.spark.range(10, numPartitions=2).selectExpr("CAST(id AS STRING) AS s")
rows = df.select(joined(pair("s"))).collect()
self.assertEqual([r[0] for r in rows], [f"{i}-{i}" for i in range(10)])
# A pass-through column makes the in-process node buffer and join input rows.
df = df.selectExpr("s", "CAST(s AS INT) AS id")
rows = df.select("id", joined(pair("s"))).collect()
self.assertEqual([tuple(r) for r in rows], [(i, f"{i}-{i}") for i in range(10)])
finally:
conf.set(key, previous)
def test_pipelined_worker_stopping_early_does_not_break_inprocess_input(self):
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
from pyspark.sql.functions import arrow_udf
conf = self.spark.sparkContext._jvm.org.apache.spark.SparkEnv.get().conf()
key = "spark.python.udf.pipelined.enabled"
previous = conf.get(key, "false")
conf.set(key, "true")
try:
double = inprocess_udf("long")(lambda x: pc.multiply(x, 2))
plus_one = arrow_udf(lambda x: pc.add(x, 1), "long")
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "100"}):
# The task completes after a few rows while the writer thread still pulls input.
# A filter on the UDF result keeps the limit from being pushed below the UDFs.
df = self.spark.range(0, 100000, 1, 2).selectExpr("id", "id % 7 AS k")
# A coalesced parent evaluates the in-process node on the writer thread.
coalesced = df.select(double("id").alias("a"), "k").coalesce(1)
queries = [
df.select("k", plus_one(double("id")).alias("v")).filter("v > 0").limit(1),
coalesced.select(plus_one("a").alias("v"), "k").filter("v > 0").limit(10),
]
for query in queries:
plan = query._jdf.queryExecution().executedPlan().toString()
self.assertNotIn("Limit", plan[plan.index("ArrowEvalPython") :])
for _ in range(5):
self.assertEqual(len(queries[0].collect()), 1)
self.assertEqual(len(queries[1].collect()), 10)
# Later tasks in the same executor still run correctly.
total = df.select(plus_one(double("id")).alias("v")).groupBy().sum("v")
self.assertEqual(total.first()[0], sum(2 * i + 1 for i in range(100000)))
finally:
conf.set(key, previous)
def test_computed_and_reordered_arguments_with_pass_through_columns(self):
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
concat = inprocess_udf("string")(lambda a, b: pc.binary_join_element_wise(a, b, ":"))
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "3"}):
df = self.spark.range(10, numPartitions=2).selectExpr(
"id", "CAST(id AS STRING) AS s", "repeat('x', CAST(id AS INT)) AS pad"
)
# Computed and reordered arguments, while id and pad pass through unchanged.
rows = df.select("id", "pad", concat(df.pad, df.s.substr(1, 1))).collect()
self.assertEqual(
[tuple(r) for r in rows], [(i, "x" * i, "x" * i + ":" + str(i)[0]) for i in range(10)]
)
def test_input_columns_read_back_from_arrow_are_unchanged(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
# Every input column is an argument of a type that reads back exactly, so the output's
# input columns come from Arrow.
count = inprocess_udf("long")(lambda *cols: pa.array([len(cols)] * len(cols[0])))
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "3"}):
df = self.spark.range(10, numPartitions=2).selectExpr(
"IF(id % 4 = 0, NULL, CAST(id AS STRING)) AS s",
"IF(id = 3, CAST(X'FF' AS STRING), 'ok') AS invalid_utf8",
"TIMESTAMP_MICROS(id * 1000003) AS ts",
"CAST(TIMESTAMP_MICROS(id) AS TIMESTAMP_NTZ) AS ntz",
"NAMED_STRUCT('a', id, 'b', id / 3, 'c', BINARY(CAST(id AS STRING))) AS st",
)
factory = self.spark._jvm.org.apache.spark.sql.execution.python
reads_back = factory.InProcessArrowEvalPythonEvaluatorFactory.readsBack
from_json = self.spark._jvm.org.apache.spark.sql.types.DataType.fromJson
for field in df.schema.fields:
self.assertTrue(reads_back(from_json(field.dataType.json())), field)
expected = df.collect()
query = df.select("*", count(*[df[c] for c in df.columns]).alias("n"))
# ReadBack also needs the arguments to be exactly the child's output, in order.
nodes = [query._jdf.queryExecution().executedPlan()]
evaluators = []
while nodes:
node = nodes.pop()
if node.nodeName() == "InProcessArrowEvalPython":
evaluators.append(node)
children = node.children().iterator()
while children.hasNext():
nodes.append(children.next())
self.assertEqual(len(evaluators), 1)
node = evaluators[0]
arguments = node.udfs().apply(0).children()
output = node.child().output()
self.assertEqual(arguments.size(), output.size())
for i in range(output.size()):
self.assertTrue(arguments.apply(i).semanticEquals(output.apply(i)), i)
rows = query.collect()
self.assertEqual([r[:-1] for r in rows], [tuple(r) for r in expected])
self.assertEqual({r.n for r in rows}, {len(df.columns)})
def test_input_columns_without_exact_arrow_read_back_keep_original_rows(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
# Intervals are converted for Arrow, and decimals, arrays and maps would read back only
# element by element, so these rows are buffered rather than read back.
count = inprocess_udf("long")(lambda *cols: pa.array([len(cols)] * len(cols[0])))
with self.sql_conf({"spark.sql.execution.arrow.maxRecordsPerBatch": "3"}):
df = self.spark.range(10, numPartitions=2).selectExpr(
"id",
"MAKE_INTERVAL(0, 1, 0, CAST(id AS INT), 0, 0, id / 1000) AS i",
"CAST(id AS DECIMAL(20, 3)) / 7 AS d",
"IF(id = 5, NULL, ARRAY(CAST(id AS STRING), NULL)) AS arr",
"MAP('k', IF(id % 2 = 0, NULL, id)) AS m",
)
# PySpark cannot collect intervals, so compare them as strings.
columns = ["id", "CAST(i AS STRING)", "d", "arr", "m"]
expected = df.selectExpr(*columns).collect()
result = df.select("*", count(*[df[c] for c in df.columns]).alias("n"))
rows = result.selectExpr(*columns, "n").collect()
self.assertEqual([r[:-1] for r in rows], [tuple(r) for r in expected])
self.assertEqual({r.n for r in rows}, {len(df.columns)})
def test_full_validation_can_be_disabled(self):
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
upper = inprocess_udf("string")(lambda s: pc.utf8_upper(s))
key = "spark.sql.execution.pythonUDF.inProcess.fullValidation.enabled"
for enabled in ["true", "false"]:
with self.sql_conf({key: enabled}):
rows = self.spark.range(3).selectExpr("CAST(id AS STRING) || 'a' AS s")
self.assertEqual(
[r[0] for r in rows.select(upper("s")).collect()], ["0A", "1A", "2A"]
)
def test_duplicate_names_in_udf_arguments_are_rejected(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
from pyspark.sql import functions as F
from pyspark.sql.types import LongType
size = inprocess_udf(LongType())(lambda x: pa.array([len(x)] * len(x), type=pa.int64()))
df = self.spark.range(3)
duplicate = F.struct(F.col("id"), F.col("id").cast("string").alias("id"))
with self.assertRaisesRegex(Exception, "DUPLICATED_FIELD_NAME"):
df.select(size(duplicate)).collect()
def test_grouping_aggregate_results_and_nested_calls(self):
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
from pyspark.sql import functions as F
from pyspark.sql.types import LongType
double = inprocess_udf(LongType())(lambda x: pc.multiply(x, 2))
df = self.spark.range(4).selectExpr("id % 2 AS k", "id AS v")
grouped = df.groupBy("k").agg(double("k").alias("r"))
self.assertEqual(sorted(r.r for r in grouped.collect()), [0, 2])
summed = df.groupBy("k").agg(F.sum("v").alias("s")).select(double("s"))
self.assertEqual(sorted(r[0] for r in summed.collect()), [4, 8])
constants = df.agg(F.count("*"), double(F.lit(1)))
self.assertEqual(tuple(constants.first()), (4, 2))
repeated = df.groupBy(double("k")).agg(double("k"))
self.assertEqual(sorted(tuple(r) for r in repeated.collect()), [(0, 0), (2, 2)])
nested = df.select(double("v").alias("x")).select(double("x"))
self.assertEqual(sorted(r[0] for r in nested.collect()), [0, 4, 8, 12])
def test_nondeterministic_grouping_and_ordering(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType(), deterministic=False)(lambda x: x)
df = self.spark.range(4)
self.assertEqual(len(df.groupBy(identity("id")).count().collect()), 4)
self.assertEqual([r.id for r in df.orderBy(identity("id")).collect()], list(range(4)))
def test_non_udf_predicates_run_before_udf(self):
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
from pyspark.sql import functions as F
from pyspark.sql.types import LongType
divide = inprocess_udf(LongType())(lambda n, d: pc.divide(n, d))
df = self.spark.createDataFrame([(4, 0), (4, 2)], "n long, d long")
result = df.filter(F.col("d") != 0).filter(divide("n", "d") > 1).collect()
self.assertEqual([tuple(r) for r in result], [(4, 2)])
def test_partition_filters_do_not_evaluate_udfs_on_driver(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql import functions as F
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType())(lambda x: x)
with tempfile.TemporaryDirectory() as path:
self.spark.range(4).selectExpr("id", "id % 2 AS k").write.partitionBy("k").parquet(
path, mode="overwrite"
)
df = self.spark.read.parquet(path)
partition_identity = inprocess_udf(df.schema["k"].dataType)(lambda x: x)
self.assertEqual(
sorted(r.id for r in df.filter(partition_identity("k") == 1).collect()), [1, 3]
)
filtered = df.filter((F.col("k") == 1) & (identity("id") > 1))
self.assertEqual([r.id for r in filtered.collect()], [3])
plan = filtered._jdf.queryExecution().executedPlan().toString()
self.assertRegex(plan, r"PartitionFilters: \[[^\]]*k#")
def test_nested_nullability_is_compatible(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql import functions as F
from pyspark.sql.types import ArrayType, StringType
identity = inprocess_udf(ArrayType(StringType()))(lambda x: x)
df = self.spark.createDataFrame([("a,b",), (None,)], "s string")
self.assertEqual(
[r[0] for r in df.select(identity(F.split("s", ","))).collect()], [["a", "b"], None]
)
def test_captured_spark_broadcasts_and_accumulators_are_rejected(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
broadcast = self.spark.sparkContext.broadcast(1)
accumulator = self.spark.sparkContext.accumulator(0)
try:
for value in (broadcast, accumulator):
with self.subTest(value=type(value).__name__):
with self.assertRaisesRegex(NotImplementedError, "broadcasts or accumulators"):
inprocess_udf(LongType())(lambda x: value.value)("id")
finally:
broadcast.destroy()
def test_sql_registration_fails_early(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType())(lambda x: x)
with self.assertRaisesRegex(TypeError, "INVALID_UDF_EVAL_TYPE"):
self.spark.udf.register("inprocess_identity", identity)
def test_named_arguments_partition_evaluator_and_columnar_child(self):
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
subtract = inprocess_udf(LongType())(lambda x, y: pc.subtract(x, y))
df = self.spark.range(8).selectExpr("id", "id + 1 as other").cache()
try:
df.count()
for evaluator in ("true", "false"):
with self.sql_conf(
{
"spark.sql.execution.usePartitionEvaluator": evaluator,
"spark.sql.execution.arrow.pythonUDF.columnarInput.enabled": "true",
"spark.sql.adaptive.enabled": "false",
}
):
query = df.select(subtract(y="id", x="other"))
self.assertEqual([r[0] for r in query.collect()], [1] * 8)
plan = query._jdf.queryExecution().executedPlan().toString()
self.assertIn("InProcessArrowEvalPython", plan)
nodes = [query._jdf.queryExecution().executedPlan()]
metrics = None
while nodes:
node = nodes.pop()
if node.nodeName() == "InProcessArrowEvalPython":
# A dual-mode cache scan can execute rows directly without a transition.
self.assertTrue(node.child().supportsColumnar())
self.assertFalse(node.supportsColumnar())
metrics = node.metrics()
children = node.children().iterator()
while children.hasNext():
nodes.append(children.next())
self.assertIsNotNone(metrics)
self.assertEqual(metrics.apply("pythonNumRowsReceived").value(), 8)
self.assertGreater(metrics.apply("pythonDataSent").value(), 0)
self.assertGreater(metrics.apply("pythonDataReceived").value(), 0)
# In nanoseconds, so a sub-millisecond registration or batch still counts.
processing = metrics.apply("pythonProcessingTime").value()
self.assertGreater(metrics.apply("pythonInitTime").value(), 0)
self.assertGreater(processing, 0)
self.assertGreaterEqual(metrics.apply("pythonTotalTime").value(), processing)
finally:
df.unpersist()
def test_nonroot_limit_and_offset_preserve_order(self):
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
identity = inprocess_udf(LongType())(lambda x: x)
with self.sql_conf(
{
"spark.sql.execution.topKSortFallbackThreshold": "1",
"spark.sql.shuffle.partitions": "4",
"spark.sql.adaptive.enabled": "false",
}
):
df = self.spark.range(100, numPartitions=4).orderBy("id")
mapped = df.select(identity("id").alias("v"))
for offset in (0, 7):
query = mapped.offset(offset).limit(10).distinct()
self.assertEqual(
sorted(r.v for r in query.collect()), list(range(offset, offset + 10))
)
self.assertIn("Sort [id#", query._jdf.queryExecution().executedPlan().toString())
def test_map_entries_offset_and_custom_names(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType, MapType, StringType
def make_map(x):
map_type = pa.map_(pa.field("k", pa.string(), False), pa.field("v", pa.int64()))
entries = pa.StructArray.from_arrays(
[pa.array(["hidden", "a", "b", "c"]), pa.array([None, 1, 2, 3])],
fields=[map_type.key_field, map_type.item_field],
)
offsets = pa.array([0, 1, 3], type=pa.int32())
return pa.Array.from_buffers(
map_type, 2, [None, offsets.buffers()[1]], children=[entries.slice(1)]
)
udf = inprocess_udf(MapType(StringType(), LongType()))(make_map)
rows = self.spark.range(2, numPartitions=1).select(udf("id")).collect()
self.assertEqual([r[0] for r in rows], [{"a": 1}, {"b": 2, "c": 3}])
def test_hidden_struct_nulls(self):
import pyarrow as pa
import pyarrow.compute as pc
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import IntegerType, StructField, StructType
@inprocess_udf(StructType([StructField("len", IntegerType(), False)]))
def lengths(s):
return pa.StructArray.from_arrays(
[pc.utf8_length(s)], names=["len"], mask=pc.is_null(s)
)
df = self.spark.createDataFrame([(None,), ("a",)], "s string")
self.assertEqual([r[0] for r in df.select(lengths("s")).collect()], [None, (1,)])
def test_large_closure_uses_bulk_transfer(self):
import pyarrow as pa
from pyspark.inprocess import inprocess_udf
from pyspark.sql.types import LongType
model = bytes(range(256)) * 16384
@inprocess_udf(LongType())
def lookup(x):
return pa.array([model[i % len(model)] for i in x.to_pylist()], type=pa.int64())
df = self.spark.range(256, numPartitions=4)
self.assertEqual([r[0] for r in df.select(lookup("id")).collect()], list(range(256)))
if __name__ == "__main__":
from pyspark.testing import main
main()