| # | |
| # 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. | |
| # | |
| """Arrow CDI contract tests that do not need a Spark JVM or JEP.""" | |
| import sys | |
| import threading | |
| import unittest | |
| import weakref | |
| from importlib.util import find_spec | |
| from unittest.mock import MagicMock, patch | |
| from pyspark import cloudpickle | |
| from pyspark.sql.types import ( | |
| ArrayType, | |
| BooleanType, | |
| DecimalType, | |
| FloatType, | |
| IntegerType, | |
| LongType, | |
| StringType, | |
| StructField, | |
| StructType, | |
| TimestampType, | |
| ) | |
| from pyspark.testing.utils import have_pyarrow | |
| _have_arrow_cdi = have_pyarrow and find_spec("cffi") is not None | |
| if _have_arrow_cdi: | |
| from pyarrow.cffi import ffi | |
| if have_pyarrow: | |
| import pyarrow as pa | |
| from pyspark.inprocess.runtime import ( | |
| _canonical_type, | |
| _has_offsets_buffers, | |
| _inprocess_invoke, | |
| _inprocess_register, | |
| _inprocess_release, | |
| _nullable_type, | |
| _Registration, | |
| _results, | |
| _strings_as_binary, | |
| _udfs, | |
| _validate_result, | |
| ) | |
| from pyspark.inprocess.udf import inprocess_udf | |
| def validate(result, expected_rows, expected_type, **options): | |
| """Validates as an invocation does, with the key that registration computes.""" | |
| return _validate_result( | |
| result, expected_rows, expected_type, _nullable_type(expected_type), **options | |
| ) | |
| @unittest.skipUnless(have_pyarrow, "In-process runtime tests require PyArrow") | |
| class InProcessRuntimeTests(unittest.TestCase): | |
| """Tests that need no Arrow C Data Interface, so they also run without cffi.""" | |
| def test_full_validation_precedes_normalization_and_export(self): | |
| # Model Arrow rejecting an invalid result without handing malformed native buffers | |
| # to either runtime. A validation failure must prevent all subsequent buffer access. | |
| result = MagicMock(spec=pa.Array) | |
| result.type = pa.string() | |
| result.__len__.return_value = 2 | |
| result.validate.side_effect = pa.ArrowInvalid("invalid result buffers") | |
| with patch("pyspark.inprocess.runtime._with_schema") as normalize: | |
| with self.assertRaisesRegex(pa.ArrowInvalid, "invalid result buffers"): | |
| validate(result, 2, pa.string()) | |
| result.validate.assert_called_once_with() | |
| normalize.assert_not_called() | |
| result.buffers.assert_not_called() | |
| # The full validation, of the binary view, also precedes every later step. | |
| result = MagicMock(spec=pa.Array) | |
| result.type = pa.string() | |
| result.__len__.return_value = 2 | |
| result.view.return_value.validate.side_effect = pa.ArrowInvalid("invalid offsets") | |
| with ( | |
| patch("pyspark.inprocess.runtime._repair_offsets") as repair, | |
| patch("pyspark.inprocess.runtime._with_schema") as normalize, | |
| ): | |
| with self.assertRaisesRegex(pa.ArrowInvalid, "invalid offsets"): | |
| validate(result, 2, pa.string()) | |
| result.validate.assert_called_once_with() | |
| result.view.assert_called_once_with(pa.binary()) | |
| result.view.return_value.validate.assert_called_once_with(full=True) | |
| repair.assert_not_called() | |
| normalize.assert_not_called() | |
| result.buffers.assert_not_called() | |
| def test_full_validation_accepts_invalid_utf8_like_spark_strings(self): | |
| # Spark strings may hold invalid UTF-8, e.g. CAST(X'FF' AS STRING); workers accept them. | |
| strings = pa.array([b"\xff", None, b"ok"], pa.binary()).view(pa.string()) | |
| values = [ | |
| strings, | |
| pa.StructArray.from_arrays([strings], names=["s"]), | |
| pa.ListArray.from_arrays(pa.array([0, 1, 3], pa.int32()), strings), | |
| pa.MapArray.from_arrays( | |
| pa.array([0, 1, 3], pa.int32()), pa.array(["a", "b", "c"]), strings | |
| ), | |
| ] | |
| for value in values: | |
| with self.subTest(type=value.type): | |
| result = validate(value, len(value), value.type) | |
| self.assertEqual( | |
| _strings_as_binary(result).to_pylist(), | |
| _strings_as_binary(value).to_pylist(), | |
| ) | |
| def test_full_validation_ignores_nullability_and_null_type_lengths(self): | |
| # Spark's StructWriter writes a null child under each null struct row. | |
| fields = [pa.field("a", pa.int32(), nullable=False), pa.field("s", pa.string())] | |
| hidden = pa.StructArray.from_arrays( | |
| [pa.array([1, None], pa.int32()), pa.array(["x", None])], | |
| fields=fields, | |
| mask=pa.array([False, True]), | |
| ) | |
| nulls = pa.array([[("k", None)], None], pa.map_(pa.string(), pa.null())) | |
| nested = pa.array( | |
| [{"l": [None, None], "s": "x"}], | |
| pa.struct([("l", pa.list_(pa.null())), ("s", pa.string())]), | |
| ) | |
| for value in [hidden, nulls, nested]: | |
| with self.subTest(type=value.type): | |
| result = validate(value, len(value), value.type) | |
| self.assertEqual(result.to_pylist(), value.to_pylist()) | |
| def test_full_validation_rejects_invalid_interior_string_offsets(self): | |
| offsets = pa.array([0, 5, 2], pa.int32()).buffers()[1] | |
| value = pa.Array.from_buffers(pa.string(), 2, [None, offsets, pa.py_buffer(b"hello")]) | |
| for result in [value, pa.StructArray.from_arrays([value], names=["s"])]: | |
| with self.subTest(type=result.type): | |
| with self.assertRaisesRegex(pa.ArrowInvalid, "non-monotonic offset"): | |
| validate(result, 2, result.type) | |
| def test_sorted_map_metadata_is_normalized_including_nested_maps(self): | |
| sorted_type = pa.map_(pa.string(), pa.int64(), keys_sorted=True) | |
| declared = pa.map_(pa.string(), pa.int64()) | |
| values = [[("a", 1), ("b", 2)], None, []] | |
| for actual_type, expected_type, data in [ | |
| (sorted_type, declared, values), | |
| (pa.list_(sorted_type), pa.list_(declared), [values]), | |
| (pa.struct([("m", sorted_type)]), pa.struct([("m", declared)]), [{"m": values[0]}]), | |
| ]: | |
| with self.subTest(actual_type=actual_type): | |
| value = pa.array(data, type=actual_type) | |
| result = validate(value, len(value), expected_type) | |
| self.assertEqual(result.type, expected_type) | |
| self.assertEqual(result.to_pylist(), value.to_pylist()) | |
| self.assertEqual( | |
| [b.address if b is not None else None for b in result.buffers()], | |
| [b.address if b is not None else None for b in value.buffers()], | |
| ) | |
| def test_session_dependent_nested_string_and_binary_widths(self): | |
| # Declare the field order: newer PyArrow versions sort inferred struct fields. | |
| value = pa.array( | |
| [{"s": ["hello", None], "b": b"data"}, None], | |
| pa.struct([("s", pa.list_(pa.string())), ("b", pa.binary())]), | |
| ) | |
| expected = pa.struct( | |
| [pa.field("s", pa.list_(pa.large_string())), pa.field("b", pa.large_binary())] | |
| ) | |
| result = validate(value, 2, expected) | |
| self.assertEqual(result.type, expected) | |
| self.assertEqual(result.to_pylist(), value.to_pylist()) | |
| self.assertEqual(validate(result, 2, value.type).to_pylist(), value.to_pylist()) | |
| def test_empty_map_does_not_read_offsets(self): | |
| from pyspark.inprocess.runtime import _null_checker | |
| expected = pa.map_(pa.string(), pa.field("value", pa.int64(), nullable=False)) | |
| empty = pa.array([], type=expected) | |
| class EmptyMap: | |
| values = empty.values | |
| null_count = 0 | |
| def __len__(self): | |
| return 0 | |
| @property | |
| def offsets(self): | |
| raise AssertionError("An empty map must not read its offsets buffer") | |
| _null_checker(expected)(EmptyMap()) | |
| nested = pa.array([[], None, []], type=pa.list_(expected)) | |
| self.assertEqual(validate(nested, 3, nested.type), nested) | |
| def test_zero_length_levels_get_offsets_buffers(self): | |
| def strings(offsets): | |
| return pa.Array.from_buffers(pa.string(), 0, [None, offsets, pa.py_buffer(b"")]) | |
| for offsets in [None, pa.py_buffer(b"")]: | |
| empty = strings(offsets) | |
| lists = pa.ListArray.from_arrays(pa.array([0] * 5, pa.int32()), empty) | |
| entries = pa.array([], pa.int64()) | |
| cases = [ | |
| (empty, 0), | |
| # A slice offset sends this through concatenation, which crashed on it. | |
| (lists.slice(1), 3), | |
| (pa.StructArray.from_arrays([lists], names=["a"]).slice(1), 3), | |
| (pa.MapArray.from_arrays(pa.array([0, 0, 0], pa.int32()), empty, entries), 2), | |
| (pa.DictionaryArray.from_arrays(pa.array([None, None], pa.int32()), empty), 2), | |
| ] | |
| for value, rows in cases: | |
| with self.subTest(type=value.type, offsets=offsets): | |
| expected = _canonical_type(value.type) | |
| result = validate(value, rows, expected) | |
| self.assertTrue(_has_offsets_buffers(result)) | |
| self.assertEqual(result.to_pylist(), value.to_pylist()) | |
| self.assertEqual(result.type, expected) | |
| def test_equivalent_representations_are_cast_to_the_declared_type(self): | |
| # A value longer than 12 bytes gives views a variadic data buffer. | |
| strings = pa.array(["a", None, "a value longer than twelve bytes"]) | |
| cases = [ | |
| (pa.array([[1], None, [2, 3]], pa.large_list(pa.int64())), pa.list_(pa.int64())), | |
| (pa.array([[1, 2], None, [3, 4]], pa.list_(pa.int64(), 2)), pa.list_(pa.int64())), | |
| (strings.cast(pa.string_view()), pa.string()), | |
| (strings.cast(pa.binary()).cast(pa.binary_view()), pa.binary()), | |
| (pa.array([b"ab", None, b"cd"], pa.binary(2)), pa.binary()), | |
| (strings.dictionary_encode(), pa.string()), | |
| ( | |
| pa.array( | |
| [{"a": [["x"]]}, None, {"a": None}], | |
| pa.struct([("a", pa.large_list(pa.list_(pa.string_view())))]), | |
| ), | |
| pa.struct([("a", pa.list_(pa.list_(pa.string())))]), | |
| ), | |
| ] | |
| for value, expected in cases: | |
| with self.subTest(type=value.type): | |
| result = validate(value, 3, expected) | |
| self.assertEqual(result.type, expected) | |
| self.assertEqual(result.to_pylist(), value.to_pylist()) | |
| def test_other_representation_differences_are_rejected(self): | |
| cases = [ | |
| (pa.array([[1], [2, 3]], pa.list_view(pa.int64())), pa.list_(pa.int64())), | |
| (pa.array([1, 2], pa.int32()).dictionary_encode(), pa.int64()), | |
| (pa.array([[1], [2, 3]], pa.large_list(pa.int32())), pa.list_(pa.int64())), | |
| ] | |
| for value, expected in cases: | |
| with self.subTest(type=value.type): | |
| with self.assertRaisesRegex(TypeError, "expected"): | |
| validate(value, 2, expected) | |
| def test_zero_argument_udf_is_rejected(self): | |
| with self.assertRaisesRegex(ValueError, "0-arg"): | |
| inprocess_udf(LongType())(lambda: pa.array([7])) | |
| def test_slices_are_normalized_including_nested_child_offsets(self): | |
| for value in ( | |
| pa.array([9, 1, None, 3]).slice(1), | |
| pa.array(["discard", "one", None, "three"]).slice(1), | |
| pa.array([[9], [1], None, [3]]).slice(1), | |
| pa.StructArray.from_arrays([pa.array([9, 1, None, 3]).slice(1)], names=["x"]), | |
| ): | |
| with self.subTest(data_type=value.type): | |
| normalized = validate(value, 3, value.type) | |
| self.assertEqual(normalized.offset, 0) | |
| self.assertEqual(normalized.to_pylist(), value.to_pylist()) | |
| if pa.types.is_struct(value.type): | |
| self.assertEqual(normalized.field(0).offset, 0) | |
| def test_nested_nullability_accepts_compatible_values(self): | |
| nullable = pa.list_(pa.field("element", pa.string(), nullable=True)) | |
| required = pa.list_(pa.field("element", pa.string(), nullable=False)) | |
| for source, expected in ((nullable, required), (required, nullable)): | |
| value = pa.array([["a"], None, []], type=source) | |
| result = validate(value, 3, expected) | |
| self.assertEqual(result.type, expected) | |
| self.assertEqual(result.to_pylist(), value.to_pylist()) | |
| with self.assertRaisesRegex(ValueError, "non-nullable"): | |
| validate(pa.array([[None]], type=nullable), 1, required) | |
| def test_null_struct_parents_do_not_violate_child_nullability(self): | |
| import pyarrow.compute as pc | |
| expected = pa.struct([pa.field("len", pa.int32(), nullable=False)]) | |
| strings = pa.array([None, "a"]) | |
| original = pa.StructArray.from_arrays( | |
| [pc.utf8_length(strings)], names=["len"], mask=pc.is_null(strings) | |
| ) | |
| values = [original, pc.if_else(pc.is_valid(strings), original, None)] | |
| values.append(pc.take(original, pa.array([None, 1], type=pa.int32()))) | |
| for value in values: | |
| with self.subTest(value=value): | |
| self.assertEqual(value.field(0).null_count, 1) | |
| result = validate(value, 2, expected) | |
| self.assertEqual(result.to_pylist(), [None, {"len": 1}]) | |
| self.assertEqual(result.type, expected) | |
| visible_null = pa.StructArray.from_arrays([pa.array([None], pa.int32())], names=["len"]) | |
| with self.assertRaisesRegex(ValueError, "non-nullable"): | |
| validate(visible_null, 1, expected) | |
| def test_sliced_map_entries_are_normalized(self): | |
| map_type = pa.map_(pa.string(), 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()) | |
| value = pa.Array.from_buffers( | |
| map_type, 2, [None, offsets.buffers()[1]], children=[entries.slice(1)] | |
| ) | |
| self.assertEqual(value.offset, 0) | |
| self.assertEqual(value.values.offset, 1) | |
| result = validate(value, 2, map_type) | |
| self.assertEqual(result.values.offset, 0) | |
| self.assertEqual(result.to_pylist(), [[("a", 1)], [("b", 2), ("c", 3)]]) | |
| def test_map_field_names_and_nested_metadata_are_normalized(self): | |
| value = pa.array( | |
| [[("a", 1)]], | |
| type=pa.map_(pa.field("k", pa.string(), False), pa.field("v", pa.int64())), | |
| ) | |
| expected = pa.map_(pa.string(), pa.int64()) | |
| result = validate(value, 1, expected) | |
| self.assertEqual(result.type.key_field.name, "key") | |
| self.assertEqual(result.type.item_field.name, "value") | |
| expected_struct = pa.struct([pa.field("x", pa.int64(), metadata={b"type": b"required"})]) | |
| result = validate(pa.array([{"x": 1}]), 1, expected_struct) | |
| self.assertEqual(result.type[0].metadata, {b"type": b"required"}) | |
| def test_map_nullability_and_sliced_results(self): | |
| nullable = pa.map_(pa.string(), pa.field("value", pa.int64())) | |
| required = pa.map_(pa.string(), pa.field("value", pa.int64(), nullable=False)) | |
| value = pa.array([[("discard", 0)], [("a", 1)], None], type=nullable).slice(1) | |
| result = validate(value, 2, required) | |
| self.assertEqual(result.offset, 0) | |
| self.assertEqual(result.type, required) | |
| self.assertEqual(result.to_pylist(), [[("a", 1)], None]) | |
| with self.assertRaisesRegex(ValueError, "non-nullable"): | |
| validate(pa.array([[("a", None)]], type=nullable), 1, required) | |
| def test_null_checks_skip_nullable_subtrees_and_null_free_parents(self): | |
| arrays = [ | |
| pa.array([{"x": [1, None]}, {"x": None}, None]), | |
| pa.array([{"x": 1}, {"x": 2}], type=pa.struct([pa.field("x", pa.int64(), False)])), | |
| pa.array([[("a", 1)], [("b", 2)]], type=pa.map_(pa.string(), pa.int64())), | |
| ] | |
| for array in arrays: | |
| with ( | |
| self.subTest(type=array.type), | |
| patch("pyspark.inprocess.runtime.pc.filter") as filtered, | |
| patch("pyspark.inprocess.runtime.pa.concat_arrays") as concat, | |
| ): | |
| result = validate(array, len(array), array.type) | |
| self.assertEqual(result, array) | |
| filtered.assert_not_called() | |
| concat.assert_not_called() | |
| def test_null_checks_still_validate_required_descendants(self): | |
| required = pa.struct([pa.field("x", pa.list_(pa.field("element", pa.int64(), False)))]) | |
| with self.assertRaisesRegex(ValueError, "non-nullable"): | |
| validate(pa.array([{"x": [None]}, None], type=required), 2, required) | |
| hidden = pa.array([{"x": None}, None], type=required) | |
| self.assertEqual(validate(hidden, 2, required), hidden) | |
| def test_map_entries_offset_respects_required_values(self): | |
| source = pa.map_(pa.string(), pa.int64()) | |
| expected = pa.map_(pa.string(), pa.field("value", pa.int64(), False)) | |
| for data in ([0, 1, 2, None], [None, 1, 2, 3]): | |
| entries = pa.StructArray.from_arrays( | |
| [pa.array(["hidden", "a", "b", "c"]), pa.array(data)], | |
| fields=[source.key_field, source.item_field], | |
| ) | |
| offsets = pa.array([0, 1, 3], pa.int32()).buffers()[1] | |
| value = pa.Array.from_buffers(source, 2, [None, offsets], children=[entries.slice(1)]) | |
| with self.subTest(data=data): | |
| if data[-1] is None: | |
| with self.assertRaisesRegex(ValueError, "non-nullable"): | |
| validate(value, 2, expected) | |
| else: | |
| result = validate(value, 2, expected) | |
| self.assertEqual(result.to_pylist(), [[("a", 1)], [("b", 2), ("c", 3)]]) | |
| def test_null_parents_do_not_copy_null_free_children(self): | |
| struct = pa.StructArray.from_arrays( | |
| [pa.array([b"payload", b"value"])], | |
| fields=[pa.field("value", pa.binary(), False)], | |
| mask=pa.array([True, False]), | |
| ) | |
| mapping = pa.MapArray.from_arrays( | |
| pa.array([0, 1, 2]), | |
| pa.array(["a", "b"]), | |
| pa.array([1, 2]), | |
| type=pa.map_(pa.string(), pa.field("value", pa.int64(), False)), | |
| mask=pa.array([True, False]), | |
| ) | |
| for value in (struct, mapping): | |
| with ( | |
| self.subTest(type=value.type), | |
| patch("pyspark.inprocess.runtime.pc.filter") as filtered, | |
| patch("pyspark.inprocess.runtime.pa.concat_arrays") as concat, | |
| ): | |
| result = validate(value, 2, value.type) | |
| self.assertEqual(result, value) | |
| filtered.assert_not_called() | |
| concat.assert_not_called() | |
| def test_null_check_does_not_copy_unchecked_siblings(self): | |
| import pyarrow.compute as pc | |
| value = pa.StructArray.from_arrays( | |
| [pa.array([None, 1]), pa.array([b"a" * 4096, b"b" * 4096])], | |
| fields=[pa.field("required", pa.int64(), False), pa.field("payload", pa.binary())], | |
| mask=pa.array([True, False]), | |
| ) | |
| with patch("pyspark.inprocess.runtime.pc.filter", wraps=pc.filter) as filtered: | |
| result = validate(value, 2, value.type) | |
| self.assertEqual(result, value) | |
| filtered.assert_called_once() | |
| self.assertEqual(filtered.call_args.args[0].type, pa.int64()) | |
| self.assertEqual( | |
| result.field(1).buffers()[2].address, value.field(1).buffers()[2].address | |
| ) | |
| def test_unsupported_types_fail_before_serialization(self): | |
| from pyspark.errors import PySparkNotImplementedError | |
| from pyspark.sql.types import ( | |
| CalendarIntervalType, | |
| CharType, | |
| VarcharType, | |
| YearMonthIntervalType, | |
| ) | |
| for declared, condition in ( | |
| (CalendarIntervalType(), "NOT_IMPLEMENTED"), | |
| (YearMonthIntervalType(), "NOT_IMPLEMENTED"), | |
| (ArrayType(YearMonthIntervalType()), "NOT_IMPLEMENTED"), | |
| (StructType([StructField("x", CalendarIntervalType())]), "NOT_IMPLEMENTED"), | |
| # Python UDFs reject CHAR/VARCHAR return types before their Arrow types. | |
| (CharType(5), "CHAR_VARCHAR_NOT_SUPPORTED_IN_PYTHON"), | |
| (VarcharType(5), "CHAR_VARCHAR_NOT_SUPPORTED_IN_PYTHON"), | |
| (StructType([StructField("x", CharType(5))]), "CHAR_VARCHAR_NOT_SUPPORTED_IN_PYTHON"), | |
| ): | |
| with self.subTest(declared=declared): | |
| wrapper = inprocess_udf(declared)(lambda x: x) | |
| with patch("pyspark.inprocess.udf._serialize_udf") as serialize: | |
| with self.assertRaises(PySparkNotImplementedError) as error: | |
| wrapper._serialize() | |
| self.assertEqual(error.exception.getCondition(), condition) | |
| serialize.assert_not_called() | |
| def test_duplicate_return_fields_fail_before_serialization(self): | |
| duplicate = StructType([StructField("x", LongType()), StructField("x", LongType())]) | |
| for declared in [ | |
| duplicate, | |
| ArrayType(duplicate), | |
| StructType([StructField("n", duplicate)]), | |
| ]: | |
| wrapper = inprocess_udf(declared)(lambda x: x) | |
| with patch("pyspark.inprocess.udf._serialize_udf") as serialize: | |
| with self.assertRaisesRegex(Exception, "DUPLICATED_FIELD_NAME"): | |
| wrapper._serialize() | |
| serialize.assert_not_called() | |
| def test_pyarrow_minimum_version_on_driver(self): | |
| # Executors check it once, when their interpreter starts. | |
| with patch.object(pa, "__version__", "1.0.0"): | |
| with self.assertRaisesRegex(ImportError, "PyArrow.*must be installed"): | |
| inprocess_udf(LongType())(lambda x: x) | |
| def test_callable_signatures(self): | |
| class NoArgs: | |
| def __call__(self): | |
| return pa.array([1]) | |
| with self.assertRaisesRegex(ValueError, "0-arg"): | |
| inprocess_udf(LongType())(NoArgs()) | |
| wrapper = inprocess_udf(LongType())(lambda **cols: cols["x"]) | |
| self.assertEqual(wrapper.func(x=pa.array([1])).to_pylist(), [1]) | |
| def test_connect_and_missing_context_have_categorized_errors(self): | |
| from pyspark import SparkContext | |
| from pyspark.errors import PySparkNotImplementedError, PySparkRuntimeError | |
| udf = inprocess_udf(LongType())(lambda x: x) | |
| with patch("pyspark.sql.utils.is_remote", return_value=True): | |
| with self.assertRaises(PySparkNotImplementedError) as error: | |
| udf("id") | |
| self.assertEqual(error.exception.getCondition(), "NOT_IMPLEMENTED") | |
| with ( | |
| patch("pyspark.sql.utils.is_remote", return_value=False), | |
| patch.object(SparkContext, "_active_spark_context", None), | |
| ): | |
| with self.assertRaises(PySparkRuntimeError) as error: | |
| udf("id") | |
| self.assertEqual(error.exception.getCondition(), "SESSION_OR_CONTEXT_NOT_EXISTS") | |
| def test_wrapper_metadata_and_return_type_validation(self): | |
| from pyspark.errors import PySparkTypeError | |
| from pyspark.util import PythonEvalType | |
| def identity(x): | |
| """Identity documentation.""" | |
| return x | |
| udf = inprocess_udf(LongType())(identity) | |
| self.assertEqual(udf.__name__, identity.__name__) | |
| self.assertEqual(udf.__doc__, identity.__doc__) | |
| # Like other Python UDFs, a callable without a name is named after its type. | |
| import functools | |
| partial = inprocess_udf(LongType())(functools.partial(identity)) | |
| self.assertEqual(partial.__name__, "partial") | |
| self.assertEqual(partial._name, "partial") | |
| self.assertIs(udf.func, identity) | |
| self.assertEqual(udf.returnType, LongType()) | |
| self.assertEqual(udf.evalType, PythonEvalType.SQL_SCALAR_ARROW_INPROCESS_UDF) | |
| self.assertTrue(udf.deterministic) | |
| self.assertIs(udf.asNondeterministic(), udf) | |
| self.assertFalse(udf.deterministic) | |
| # The keyword and the function form match udf, pandas_udf and arrow_udf. | |
| self.assertEqual(inprocess_udf(returnType=LongType())(identity).returnType, LongType()) | |
| direct = inprocess_udf(identity, LongType()) | |
| self.assertIsInstance(direct, type(udf)) | |
| self.assertEqual(direct.returnType, LongType()) | |
| self.assertFalse(inprocess_udf(identity, "long", deterministic=False).deterministic) | |
| # Misuse fails at once, not when the decorator or the UDF is called. | |
| for args, kwargs, arg_name in [ | |
| ((42,), {}, "f"), | |
| ((), {}, "returnType"), | |
| ((identity,), {}, "returnType"), | |
| ((identity, 42), {}, "returnType"), | |
| (("long",), {"deterministic": "no"}, "deterministic"), | |
| ]: | |
| with self.subTest(args=args, kwargs=kwargs): | |
| with self.assertRaises(PySparkTypeError) as error: | |
| inprocess_udf(*args, **kwargs) | |
| self.assertEqual(error.exception.getMessageParameters()["arg_name"], arg_name) | |
| # The decorator form checks the function as the function form does. | |
| for decorator in [inprocess_udf("long"), inprocess_udf(returnType=LongType())]: | |
| with self.subTest(decorator=decorator): | |
| with self.assertRaises(PySparkTypeError) as error: | |
| decorator(42) | |
| self.assertEqual(error.exception.getCondition(), "NOT_EXPECTED_TYPE") | |
| self.assertEqual(error.exception.getMessageParameters()["arg_name"], "f") | |
| with self.assertRaises(PySparkTypeError) as error: | |
| inprocess_udf(LongType(), returnType=LongType()) | |
| self.assertEqual(error.exception.getCondition(), "CANNOT_SET_TOGETHER") | |
| with self.assertRaises(TypeError): | |
| inprocess_udf(LongType(), False) | |
| # DDL decoration does not require a live Spark session. | |
| self.assertIsNone(inprocess_udf("long")(identity)._parsed_return_type) | |
| @unittest.skipUnless(_have_arrow_cdi, "Arrow CDI tests require PyArrow and cffi") | |
| class InProcessRuntimeCDITests(unittest.TestCase): | |
| def tearDown(self): | |
| _results.clear() | |
| _udfs.clear() | |
| def register(self, handle, serialized, expected=None, version=None, **options): | |
| schema = ffi.new("struct ArrowSchema*") | |
| address = int(ffi.cast("uintptr_t", schema)) | |
| field = expected if expected is not None else pa.field("result", pa.int64()) | |
| field._export_to_c(address) | |
| try: | |
| _inprocess_register( | |
| handle, serialized, address, version or "%d.%d" % sys.version_info[:2], **options | |
| ) | |
| self.assertEqual(schema.release, ffi.NULL) | |
| finally: | |
| if schema.release != ffi.NULL: | |
| schema.release(schema) | |
| def invoke(self, func, inputs, return_type, rows=None, timezone="UTC"): | |
| arrays = [ffi.new("struct ArrowArray*") for _ in inputs] | |
| schemas = [ffi.new("struct ArrowSchema*") for _ in inputs] | |
| output = ffi.new("struct ArrowArray*") | |
| output_schema = ffi.new("struct ArrowSchema*") | |
| def address(value): | |
| return int(ffi.cast("uintptr_t", value)) | |
| try: | |
| for value, array, schema in zip(inputs, arrays, schemas): | |
| value._export_to_c(address(array), address(schema)) | |
| serialized = ( | |
| func._serialize() if hasattr(func, "_serialize") else cloudpickle.dumps(func) | |
| ) | |
| from pyspark.sql.pandas.types import to_arrow_type | |
| self.register( | |
| "test", | |
| serialized, | |
| pa.field("result", to_arrow_type(return_type, timezone=timezone)), | |
| ) | |
| _inprocess_invoke( | |
| "test", | |
| [address(a) for a in arrays], | |
| [address(s) for s in schemas], | |
| address(output), | |
| address(output_schema), | |
| len(inputs[0]) if rows is None else rows, | |
| ) | |
| return pa.Array._import_from_c(address(output), address(output_schema)) | |
| finally: | |
| _inprocess_release(["test"]) | |
| for value in arrays + schemas + [output, output_schema]: | |
| if value.release != ffi.NULL: | |
| value.release(value) | |
| def test_worker_style_command_is_rejected_at_registration(self): | |
| with self.assertRaisesRegex(RuntimeError, "must contain a callable; use inprocess_udf"): | |
| self.register("worker", cloudpickle.dumps((lambda x: x, LongType()))) | |
| self.assertNotIn("worker", _udfs) | |
| def test_validation_errors_do_not_capture_unvalidated_result_locals(self): | |
| def wrong_length(x): | |
| return pa.array([1, 2, 3]) | |
| def failing(x): | |
| raise ValueError("user error") | |
| for func, expected_locals in [(wrong_length, False), (failing, True)]: | |
| with self.subTest(func=func.__name__): | |
| self.register("locals", cloudpickle.dumps(func), traceback_with_locals=True) | |
| array = ffi.new("struct ArrowArray*") | |
| schema = ffi.new("struct ArrowSchema*") | |
| pa.array([1, 2], pa.int64())._export_to_c( | |
| int(ffi.cast("uintptr_t", array)), int(ffi.cast("uintptr_t", schema)) | |
| ) | |
| with patch( | |
| "pyspark.inprocess.runtime._format_exception", return_value="formatted" | |
| ) as format_exception: | |
| with self.assertRaisesRegex(RuntimeError, "formatted"): | |
| _inprocess_invoke( | |
| "locals", | |
| [int(ffi.cast("uintptr_t", array))], | |
| [int(ffi.cast("uintptr_t", schema))], | |
| 0, | |
| 0, | |
| 2, | |
| ) | |
| self.assertEqual(format_exception.call_args.args[3], expected_locals) | |
| _inprocess_release(["locals"]) | |
| def test_full_validation_can_be_disabled_per_registration(self): | |
| offsets = pa.array([0, 5, 2], pa.int32()).buffers()[1] | |
| value = pa.Array.from_buffers(pa.string(), 2, [None, offsets, pa.py_buffer(b"hello")]) | |
| # Only constant-time checks remain, which do not inspect interior offsets. | |
| self.assertEqual(len(validate(value, 2, value.type, full_validation=False)), 2) | |
| self.register("full", cloudpickle.dumps(lambda x: x)) | |
| self.register("constant", cloudpickle.dumps(lambda x: x), full_validation=False) | |
| self.assertTrue(_udfs["full"].full_validation) | |
| self.assertFalse(_udfs["constant"].full_validation) | |
| def test_identity_retains_buffers_and_nulls(self): | |
| value = pa.array([1, None, 3], type=pa.int64()) | |
| result = self.invoke(lambda x: x, [value], LongType()) | |
| self.assertEqual(result, value) | |
| self.assertEqual(result.buffers()[1].address, value.buffers()[1].address) | |
| def test_wrong_length(self): | |
| for delta in (-1, 1): | |
| with ( | |
| self.subTest(delta=delta), | |
| self.assertRaisesRegex(RuntimeError, "returned .* rows; expected 3"), | |
| ): | |
| self.invoke( | |
| lambda x: pa.array([1] * (len(x) + delta)), | |
| [pa.array([1, 2, 3])], | |
| LongType(), | |
| ) | |
| def test_wrong_return_object_has_traceback(self): | |
| with self.assertRaisesRegex(RuntimeError, "must return a pyarrow.Array") as error: | |
| self.invoke(lambda x: [1, 2], [pa.array([1, 2])], LongType()) | |
| self.assertIn("__INPROCESS_UDF_TRACEBACK__:", str(error.exception)) | |
| self.assertIn("Traceback", str(error.exception)) | |
| def test_wrong_declared_type(self): | |
| with self.assertRaisesRegex(RuntimeError, "expected string"): | |
| self.invoke(lambda x: x, [pa.array([1, 2])], StringType()) | |
| def test_nested_schema_mismatch(self): | |
| expected = StructType([StructField("values", ArrayType(StringType()))]) | |
| value = pa.array([{"values": [1, 2]}]) | |
| with self.assertRaisesRegex(RuntimeError, "expected struct"): | |
| self.invoke(lambda x: x, [value], expected) | |
| def test_decimal_scale_mismatch(self): | |
| from decimal import Decimal | |
| value = pa.array([Decimal("1.2")], type=pa.decimal128(10, 1)) | |
| with self.assertRaisesRegex(RuntimeError, "expected decimal128"): | |
| self.invoke(lambda x: x, [value], DecimalType(10, 2)) | |
| def test_timestamp_timezone_labels_preserve_instants_and_buffers(self): | |
| value = pa.array([0, None, 123456], type=pa.timestamp("us", tz="UTC")) | |
| for timezone in ["Etc/UTC", "America/Los_Angeles"]: | |
| result = self.invoke(lambda x: x, [value], TimestampType(), timezone=timezone) | |
| self.assertEqual(result.type, pa.timestamp("us", tz=timezone)) | |
| self.assertEqual(result.cast(pa.int64()).to_pylist(), [0, None, 123456]) | |
| self.assertEqual(result.buffers()[1].address, value.buffers()[1].address) | |
| for datatype in [pa.timestamp("us"), pa.timestamp("ms", tz="UTC")]: | |
| with self.assertRaisesRegex(TypeError, "expected"): | |
| validate(pa.array([0], type=datatype), 1, value.type) | |
| def test_exported_numpy_buffers_are_finalized_on_the_interpreter_thread(self): | |
| import numpy as np | |
| from pyspark.inprocess import runtime | |
| finalized = [] | |
| owners = [] | |
| def produce(values): | |
| array = np.arange(len(values), dtype=np.int64) | |
| owners.append(weakref.ref(array)) | |
| weakref.finalize(array, lambda: finalized.append(threading.get_ident())) | |
| return pa.array(array) | |
| _udfs["owned"] = _Registration( | |
| produce, pa.int64(), pa.int64(), lambda array: None, False, False, False, True | |
| ) | |
| for batch in range(2): | |
| array = ffi.new("struct ArrowArray*") | |
| schema = ffi.new("struct ArrowSchema*") | |
| # A normal input is imported on the simulated interpreter thread. | |
| input_array = ffi.new("struct ArrowArray*") | |
| input_schema = ffi.new("struct ArrowSchema*") | |
| pa.array([1, 2])._export_to_c( | |
| int(ffi.cast("uintptr_t", input_array)), | |
| int(ffi.cast("uintptr_t", input_schema)), | |
| ) | |
| runtime._inprocess_invoke( | |
| "owned", | |
| [int(ffi.cast("uintptr_t", input_array))], | |
| [int(ffi.cast("uintptr_t", input_schema))], | |
| int(ffi.cast("uintptr_t", array)), | |
| int(ffi.cast("uintptr_t", schema)), | |
| 2, | |
| ) | |
| def release_cdi(): | |
| array.release(array) | |
| schema.release(schema) | |
| task = threading.Thread(target=release_cdi) | |
| task.start() | |
| task.join() | |
| self.assertIsNotNone(owners[-1]()) | |
| self.assertEqual(len(finalized), batch) | |
| _inprocess_release(["owned"]) | |
| self.assertEqual(finalized, [threading.get_ident()] * 2) | |
| self.assertTrue(all(owner() is None for owner in owners)) | |
| def test_primitive_types_do_not_implicitly_cast(self): | |
| cases = [ | |
| (pa.array([1, 2]), IntegerType()), | |
| (pa.array(["1", "22"]), LongType()), | |
| (pa.array([1, 2], type=pa.timestamp("us", tz="UTC")), LongType()), | |
| (pa.array([1, 2], type=pa.date32()), IntegerType()), | |
| (pa.array([0.001, 0.0]), BooleanType()), | |
| (pa.array([1e300, 0.0]), FloatType()), | |
| ] | |
| for value, declared in cases: | |
| with self.subTest(value=value.type, declared=declared): | |
| wrapper = inprocess_udf(declared)(lambda x: x) | |
| with self.assertRaisesRegex(RuntimeError, "expected"): | |
| self.invoke(wrapper, [value], declared) | |
| def test_serialization_is_deferred_until_first_use(self): | |
| namespace = {"inprocess_udf": inprocess_udf, "LongType": LongType, "pa": pa} | |
| exec( | |
| "@inprocess_udf(LongType())\n" | |
| "def f(x): return pa.array([LOOKUP] * len(x), type=pa.int64())\n", | |
| namespace, | |
| ) | |
| wrapper = namespace["f"] | |
| namespace["LOOKUP"] = 42 | |
| self.assertEqual(self.invoke(wrapper, [pa.array([0])], LongType()).to_pylist(), [42]) | |
| # The first call creates the JVM-side function, which then keeps the command; see | |
| # test_command_is_stable_after_the_first_call in test_inprocess_udf. | |
| def test_empty_batch(self): | |
| value = pa.array([], type=pa.int64()) | |
| self.assertEqual(self.invoke(lambda x: x, [value], LongType()), value) | |
| def test_user_error_includes_traceback(self): | |
| def fail(x): | |
| raise ValueError("expected failure") | |
| with self.assertRaisesRegex(RuntimeError, "expected failure") as error: | |
| self.invoke(fail, [pa.array([1])], LongType()) | |
| self.assertIn("Traceback", str(error.exception)) | |
| def test_system_exit_is_converted_to_an_ordinary_exception(self): | |
| def fail(x): | |
| raise SystemExit(0) | |
| with self.assertRaisesRegex(RuntimeError, "SystemExit"): | |
| self.invoke(fail, [pa.array([1])], LongType()) | |
| def test_base_exception_during_deserialization_is_converted(self): | |
| def fail(): | |
| raise SystemExit(0) | |
| class FailingLoad: | |
| def __reduce__(self): | |
| return fail, () | |
| with self.assertRaisesRegex(RuntimeError, "SystemExit"): | |
| self.register("bad", cloudpickle.dumps(FailingLoad())) | |
| self.assertNotIn("bad", _udfs) | |
| def test_python_version_is_checked_before_deserialization(self): | |
| with self.assertRaisesRegex(RuntimeError, "IN_PROCESS_PYTHON_VERSION_MISMATCH") as error: | |
| self.register("bad", b"invalid pickle", version="0.0") | |
| # The message names what selects the embedded interpreter, not PYSPARK_PYTHON alone. | |
| self.assertIn("JEP built for Python 0.0", str(error.exception)) | |
| self.assertNotIn("bad", _udfs) | |
| def test_registration_is_task_scoped(self): | |
| state = [] | |
| def remember(x): | |
| state.append(x) | |
| return len(state) | |
| command = cloudpickle.dumps(remember) | |
| for handle in ("first", "second"): | |
| self.register(handle, command) | |
| self.assertEqual(_udfs["first"].func(1), 1) | |
| self.assertEqual(_udfs["first"].func(2), 2) | |
| self.assertEqual(_udfs["second"].func(3), 1) | |
| _inprocess_release(["first", "second", "unregistered"]) | |
| self.assertFalse(_udfs) | |
| def test_registration_consumes_declared_cdi_schema(self): | |
| # An independently provided field is authoritative, including nested metadata, | |
| # child names, large binary layout, and nullability. | |
| expected = pa.field( | |
| "result", | |
| pa.struct( | |
| [pa.field("payload", pa.large_binary(), nullable=False, metadata={"logical": "v"})] | |
| ), | |
| ) | |
| self.register("schema", cloudpickle.dumps(lambda x: x), expected) | |
| self.assertTrue(_udfs["schema"].expected_type.equals(expected.type, check_metadata=True)) | |
| with self.assertRaisesRegex(ValueError, "non-nullable"): | |
| _udfs["schema"].checker(pa.array([{"payload": None}], type=expected.type)) | |
| def test_exception_text_is_safe_for_jni(self): | |
| message = "failure: caf\u00e9 \u4e2d\u6587 \U0001f600 \ud800 \0 tail" | |
| def fail(x): | |
| raise ValueError(message) | |
| with self.assertRaises(RuntimeError) as error: | |
| self.invoke(fail, [pa.array([1])], LongType()) | |
| text = str(error.exception) | |
| self.assertIn("caf\u00e9 \u4e2d\u6587", text) | |
| self.assertNotIn("\0", text) | |
| self.assertIn(r"\U0001f600 \ud800 \x00 tail", text) | |
| def test_registration_traceback_policy(self): | |
| for hide in [False, True]: | |
| with self.assertRaises(RuntimeError) as error: | |
| self.register("bad", b"", version="0.0", hide_traceback=hide) | |
| self.assertEqual("Traceback" in str(error.exception), not hide) | |
| self.assertIn("IN_PROCESS_PYTHON_VERSION_MISMATCH", str(error.exception)) | |
| if __name__ == "__main__": | |
| from pyspark.testing import main | |
| main() |