diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 88141df3..e02722aa 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -4089,33 +4089,23 @@ fn check_collectable( Ok(()) } -/// Whether the binary values of a result of `schema` of `dataframe` are `bytes`, as PySpark -/// Connect's `DataFrame._get_binary_as_bytes` gives it: the conf -/// `spark.sql.execution.pyspark.binaryAsBytes`, true unless it is set to another value than -/// `true`. The conf is read only for a result with a binary value, the only one it changes. -fn binary_as_bytes( - py: Python<'_>, - dataframe: &DataFrame, - schema: Option<&spark_connect::types::DataType>, -) -> PyResult { - let binary = - |t: &spark_connect::types::DataType| matches!(t, spark_connect::types::DataType::Binary); - if !schema.is_some_and(|schema| has_type(schema, &binary)) { - return Ok(true); - } - let value = py - .detach(|| { - dataframe - .spark_session() - .conf() - .get_with_default("spark.sql.execution.pyspark.binaryAsBytes", Some("true")) - }) - .to_pyerr()?; +/// The binary result setting, as PySpark's `DataFrame._get_binary_as_bytes` reads it, +/// including when the server sends only an Arrow schema or the result has no binary fields. +/// Collect reads it after execution; a local iterator snapshots it before its first fetch. +fn read_binary_as_bytes(py: Python<'_>, dataframe: &DataFrame) -> PyResult { + let value = detached(py, || { + dataframe + .spark_session() + .conf() + .get_with_default("spark.sql.execution.pyspark.binaryAsBytes", Some("true")) + .map_err(Box::new) + })?; Ok(value.is_some_and(|value| value.eq_ignore_ascii_case("true"))) } /// The collected `rows` of a result of `schema` of `dataframe`, whose values of user-defined -/// types are deserialized when first read. +/// types are deserialized when first read. Reads the binary result setting after execution +/// and applies it to binary values, including when only an Arrow schema is available. fn rows_to_py( py: Python<'_>, dataframe: &DataFrame, @@ -4123,9 +4113,9 @@ fn rows_to_py( schema: Option<&spark_connect::types::DataType>, ) -> PyResult> { check_arrow_schema(py, schema)?; + let binary_as_bytes = read_binary_as_bytes(py, dataframe)?; check_collectable(py, schema)?; let converters = udt_field_converters(py, schema)?; - let binary_as_bytes = binary_as_bytes(py, dataframe, schema)?; Ok(rows .into_iter() .map(|row| PyRow::with_converters(row, converters.clone(), binary_as_bytes)) @@ -4146,8 +4136,8 @@ pub struct PyLocalRowIterator { /// The converters of the fields of user-defined types of the rows, once the schema of the /// result is known: with the first row. converters: Option>, - /// Whether the binary values of the rows are `bytes`, known with the converters. - binary_as_bytes: bool, + /// Snapshotted before the first fetch, as the body of PySpark's generator does. + binary_as_bytes: Option, } impl PyLocalRowIterator { @@ -4158,7 +4148,7 @@ impl PyLocalRowIterator { iterator: None, finished: false, converters: None, - binary_as_bytes: true, + binary_as_bytes: None, } } @@ -4172,6 +4162,38 @@ impl PyLocalRowIterator { } self.iterator.as_mut()?.next() } + + /// Prepare and fetch one Python row; the public wrapper closes on exhaustion or error. + fn next_python_row(&mut self, py: Python<'_>) -> PyResult> { + let binary_as_bytes = match self.binary_as_bytes { + Some(value) => value, + None => { + let value = read_binary_as_bytes(py, &self.dataframe)?; + self.binary_as_bytes = Some(value); + value + } + }; + let next = py.detach(|| self.next_row()); + if self.converters.is_none() { + if let Some(schema) = self.iterator.as_ref().and_then(LocalRowIterator::schema) { + // The result is known once its first message arrives, possibly with no rows. + check_collectable(py, Some(&schema))?; + self.converters = Some(udt_field_converters(py, Some(&schema))?); + } + } + match next { + Some(Ok(row)) => { + let converters = self.converters.clone().flatten(); + Ok(Some(PyRow::with_converters( + row, + converters, + binary_as_bytes, + ))) + } + Some(Err(error)) => Err(crate::errors::spark_error_to_py_exception(error)), + None => Ok(None), + } + } } #[pymethods] @@ -4191,42 +4213,17 @@ impl PyLocalRowIterator { /// Fetch the next row from the iterator. /// - /// Releases the GIL during the blocking fetch to allow other threads to run. + /// The first call snapshots the binary result setting before starting execution. + /// Both the Config RPC and blocking row fetch release the GIL. fn __next__(&mut self, py: Python<'_>) -> PyResult> { if self.finished { return Ok(None); } - // Release the GIL while fetching from the Rust iterator (which may block on network I/O). - let next = py.detach(|| self.next_row()); - if self.converters.is_none() { - if let Some(schema) = self.iterator.as_ref().and_then(LocalRowIterator::schema) { - // The result is known once its first message arrives, possibly with no rows. - if let Err(error) = check_collectable(py, Some(&schema)) { - self.finished = true; - return Err(error); - } - let converters = udt_field_converters(py, Some(&schema))?; - self.binary_as_bytes = binary_as_bytes(py, &self.dataframe, Some(&schema))?; - self.converters = Some(converters); - } - } - match next { - Some(Ok(row)) => { - let converters = self.converters.clone().flatten(); - Ok(Some(PyRow::with_converters( - row, - converters, - self.binary_as_bytes, - ))) - } - // As PySpark raises the error of the execution, which ends the iteration. - Some(Err(error)) => { - self.finished = true; - Err(crate::errors::spark_error_to_py_exception(error)) - } - None => { - self.finished = true; - Ok(None) + match self.next_python_row(py) { + Ok(Some(row)) => Ok(Some(row)), + result => { + self.close(py); + result } } } diff --git a/crates/pyspark-rs/src/functions.rs b/crates/pyspark-rs/src/functions.rs index 74da2f4f..14ac379f 100644 --- a/crates/pyspark-rs/src/functions.rs +++ b/crates/pyspark-rs/src/functions.rs @@ -333,9 +333,11 @@ fn numpy_ndarray_ndim(obj: &Bound<'_, PyAny>) -> PyResult> { /// The literal of a NumPy scalar, typed as PySpark's `LiteralExpression._infer_type` types it: a /// bool, a signed int of 8 to 64 bits (tinyint to bigint), or a float32 or float64. A str or bytes /// scalar is `None`, for the string and binary conversions to take. Any other NumPy scalar (an -/// unsigned int, float16, longdouble, complex, datetime64 or timedelta64) is -/// UNSUPPORTED_DATA_TYPE, named by its Python type. NumPy is looked up in `sys.modules`, as a -/// NumPy scalar exists only when NumPy is imported. +/// unsigned int, float16, complex, datetime64 or timedelta64) is UNSUPPORTED_DATA_TYPE, +/// named by its Python type. A supported dtype with a scalar class outside PySpark's accepted +/// classes raises a bare `AssertionError`, such as a distinct `longlong`, or an eight-byte +/// `longdouble` on macOS arm64 and Windows. Wider `longdouble` dtypes are unsupported. +/// NumPy is looked up in `sys.modules`, as a NumPy scalar exists only when NumPy is imported. fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult> { let py = obj.py(); let modules = py.import("sys")?.getattr("modules")?; @@ -352,6 +354,36 @@ fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult &["int8"], + ("i", 2) => &["int16", "int8"], + ("i", 4) => &["int32", "int16", "int8"], + ("i", 8) => &["int64", "int32", "int16", "int8"], + ("f", 4) => &["float32"], + ("f", 8) => &["float64", "float32"], + _ => &[], + }; + if !scalar_classes.is_empty() { + let mut accepted = match kind.as_str() { + "i" => obj.is_instance_of::(), + "f" => obj.is_instance_of::(), + _ => false, + }; + if !accepted { + for class in scalar_classes { + if obj.is_instance(&numpy.getattr(*class)?)? { + accepted = true; + break; + } + } + } + if !accepted { + return Err(pyo3::exceptions::PyAssertionError::new_err(())); + } + } let literal = match (kind.as_str(), itemsize) { ("U", _) | ("S", _) => return Ok(None), ("i", 1) => LiteralExpression::Byte(obj.extract::()? as i32), diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index bafdebba..166f2ad2 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -199,10 +199,11 @@ impl Iterator for LocalRowIterator { } } -/// Spawn a background task that drains the response stream and forwards decoded -/// batches over a bounded (capacity-1) channel, keeping one batch buffered ahead -/// of the consumer. Execution metrics are recorded on the session once the -/// stream ends. Backs `to_local_iterator(prefetch_partitions = true)`. +/// Spawn a background task that forwards decoded response batches over a bounded +/// (capacity-1) channel, keeping one batch buffered ahead of the consumer. +/// Closing the receiver stops the task even while it awaits a response; dropping +/// the stream releases the execution. Metrics received so far are recorded on the +/// session when the task stops. Backs `to_local_iterator(prefetch_partitions = true)`. fn spawn_prefetch( session: SparkSession, mut stream: ExecutionStream, @@ -212,7 +213,13 @@ fn spawn_prefetch( get_runtime().spawn(async move { let mut execution_info = ExecutionInfo::default(); loop { - match stream.message().await { + // Closing the consumer must release the execution even while the server has + // not produced another response. Dropping this task's stream releases it. + let response = tokio::select! { + _ = tx.closed() => break, + response = stream.message() => response, + }; + match response { Ok(Some(mut resp)) => { capture_execution(&mut resp, &mut execution_info, &session); if let Err(e) = keep_schema(&mut resp, &schema) { diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index 0d4beee0..38f49653 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -254,6 +254,10 @@ class Service: response_gate = Event() # With `block_execute` set, ExecutePlan waits for the gate too, as AnalyzePlan does. block_execute = False + # Hold r2 after the first Arrow batch so cleanup tests cannot confuse an + # execution released by close with one released by normal completion. + block_result_complete = False + result_complete_operations = [] # Each AnalyzePlanRequest, recorded. analyze_requests = [] @@ -265,10 +269,14 @@ def reset(self): self.data_type = struct_schema() self.failure = None self.block_response = False - self.response_gate.clear() + # An old ExecutePlan generator keeps its own gate across test resets. + self.response_gate = Event() self.block_execute = False + self.block_result_complete = False + self.result_complete_operations = [] self.analyze_requests = [] self.execute_requests = [] + self.release_execute_requests = [] self.execute_payload = b"" self.execute_failure = None self.execute_error = None @@ -280,6 +288,10 @@ def reset(self): self.unset_confs = [] self.confs = {} self.read_confs = [] + self.config_requests = [] + self.config_failure = None + self.block_config = False + self.rpc_events = [] self.set_confs = [] self.clone_requests = [] # The DDL of the JSON of each JsonToDDL. @@ -331,6 +343,7 @@ def analyze(self, request, context): # with the status code name and message of a tuple. Responses echo the request's # operation id, as the server does. execute_requests = [] + release_execute_requests = [] execute_payload = b"" execute_failure = None # With `execute_error` set to (message, reason, metadata), ExecutePlan fails as the @@ -362,12 +375,22 @@ def analyze(self, request, context): "spark.sql.session.localRelationBatchOfChunksSizeBytes": "1073741824", } read_confs = [] + config_requests = [] + config_failure = None + block_config = False + rpc_events = [] set_confs = [] def execute(self, request, context): - self.execute_calls += 1 + # Capture these before yielding: an earlier call can resume after reset. + response_gate = self.response_gate + block_result_complete = self.block_result_complete + result_complete_operations = self.result_complete_operations + # Counter waiters on another thread must see both records when ready. self.execute_requests.append(request) - if self.block_execute and not self.response_gate.wait(5): + self.rpc_events.append("execute") + self.execute_calls += 1 + if self.block_execute and not response_gate.wait(5): context.abort( grpc.StatusCode.DEADLINE_EXCEEDED, "Python thread could not run while the statement ran", @@ -402,12 +425,22 @@ def execute(self, request, context): operation_id = parse_message(request).get(6, [b""])[0] header = message_field(12, operation_id) yield header + message_field(13, "r1") + self.execute_payload + if block_result_complete and not response_gate.wait(30): + context.abort( + grpc.StatusCode.DEADLINE_EXCEEDED, + "ResultComplete was not released by the test", + ) + result_complete_operations.append(operation_id) yield header + message_field(13, "r2") + message_field(14, b"") def warnings_fields(self): # ConfigResponse.warnings=3, of `config_warnings`. return b"".join(message_field(3, warning) for warning in self.config_warnings) + def release_execute(self, request, context): + self.release_execute_requests.append(request) + return b"" + def fetch_error_details(self, request, context): # FetchErrorDetailsRequest.error_id=3. self.fetched_error_ids.append(parse_message(request)[3][0].decode()) @@ -429,6 +462,15 @@ def clone_session(self, request, context): def config(self, request, context): # ConfigRequest.session_id=1, operation=3, Operation.get=2, Get.keys=1; # ConfigResponse.session_id=1, pairs=2, KeyValue.key=1, value=2. + self.config_requests.append(request) + self.rpc_events.append("config") + if self.block_config and not self.response_gate.wait(5): + context.abort( + grpc.StatusCode.DEADLINE_EXCEEDED, + "Python thread could not run while configuration lookup waited", + ) + if self.config_failure: + context.abort(grpc.StatusCode.INVALID_ARGUMENT, self.config_failure) request = parse_message(request) operation = parse_message(request[3][0]) if 1 in operation: @@ -450,14 +492,20 @@ def config(self, request, context): return message_field(1, request[1][0]) + self.warnings_fields() if 3 in operation: # Operation.get_with_default=3: GetWithDefault.pairs=1 (KeyValue.key=1, - # value=2), the default of each key that `confs` does not have. + # value=2), the default of each key that `confs` does not have. A missing + # default stays unset in the response. pairs = [] for pair in parse_message(operation[3][0]).get(1, []): pair = parse_message(pair) key = pair[1][0].decode() self.read_confs.append(key) - value = self.confs.get(key, pair.get(2, [b""])[0].decode()) - pairs.append(message_field(2, message_field(1, key) + message_field(2, value))) + default = pair.get(2, [None])[0] + default = None if default is None else default.decode() + value = self.confs.get(key, default) + encoded = message_field(1, key) + if value is not None: + encoded += message_field(2, value) + pairs.append(message_field(2, encoded)) return message_field(1, request[1][0]) + b"".join(pairs) if 2 not in operation: context.abort( @@ -485,7 +533,7 @@ def config(self, request, context): "ReattachExecute": grpc.unary_stream_rpc_method_handler( lambda request, context: iter([message_field(14, b"")]) ), - "ReleaseExecute": grpc.unary_unary_rpc_method_handler(lambda request, context: b""), + "ReleaseExecute": grpc.unary_unary_rpc_method_handler(service.release_execute), "FetchErrorDetails": grpc.unary_unary_rpc_method_handler(service.fetch_error_details), "Config": grpc.unary_unary_rpc_method_handler(service.config), "CloneSession": grpc.unary_unary_rpc_method_handler(service.clone_session), diff --git a/python/tests/binary_configuration_cases.py b/python/tests/binary_configuration_cases.py new file mode 100644 index 00000000..372b45b5 --- /dev/null +++ b/python/tests/binary_configuration_cases.py @@ -0,0 +1,333 @@ +# 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. +"""Binary result configuration through native actions and the mock transport.""" + +from concurrent.futures import ThreadPoolExecutor +import subprocess +import sys +import time + +import pytest + +from _mock_connect_server import data_type_proto, message_field, parse_message, varint_field +from pyspark.errors import PySparkException +from pyspark.sql import types as T + + +_BINARY_CONF = "spark.sql.execution.pyspark.binaryAsBytes" +_EAGER_ACTIONS = ["collect", "first", "head", "head_n", "take", "tail"] +_ACTIONS = _EAGER_ACTIONS + ["iterator", "prefetch"] + + +_ITERATOR_RELEASE_CHILD = r""" +import sys +import pytest +from pyspark.sql import SparkSession, types as T + +spark = SparkSession.builder.remote(sys.argv[1]).getOrCreate() +reason = sys.argv[3] +original = T._udt_field_converters + +def fail(schema): + raise LookupError("converter initialization failed") + +if reason == "converter_failure": + T._udt_field_converters = fail +iterator = spark.range(2).toLocalIterator(sys.argv[2] == "True") +try: + if reason == "converter_failure": + with pytest.raises(LookupError, match="converter initialization failed"): + next(iterator) + T._udt_field_converters = original + else: + assert next(iterator).b == b"x" + iterator.close() + with pytest.raises(StopIteration): + next(iterator) + # Keep the process and its runtime alive until the parent verifies ReleaseAll. + assert sys.stdin.readline() == "finish\n" +finally: + T._udt_field_converters = original + iterator.close() + spark.stop() +""" + + +def _results(service, batch, schema=None): + import pyarrow as pa + + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, batch.schema) as writer: + writer.write_batch(batch) + service.execute_payload = message_field( + 2, varint_field(1, batch.num_rows) + message_field(2, sink.getvalue().to_pybytes()) + ) + if schema is not None: + service.execute_payload += message_field(7, data_type_proto(schema.jsonValue())) + + +def _action(frame, method): + if method in ("first", "head"): + row = getattr(frame, method)() + return [] if row is None else [row] + if method in ("head_n", "take", "tail"): + return getattr(frame, "head" if method == "head_n" else method)(1) + if method == "collect": + return frame.collect() + return list(frame.toLocalIterator(method == "prefetch")) + + +def _wait_for(predicate, timeout=5): + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() >= deadline: + raise AssertionError("mock RPC did not arrive") + time.sleep(0.01) + + +class TestBinaryConfiguration: + @pytest.mark.parametrize("method", _ACTIONS) + @pytest.mark.parametrize("include_schema", [False, True]) + @pytest.mark.parametrize("values", [[], [None], [b"value"]]) + def test_result_setting_applies_without_a_spark_schema( + self, mock_connect, method, include_schema, values + ): + pa = pytest.importorskip("pyarrow") + spark, service = mock_connect + service.confs = {_BINARY_CONF: "false"} + batch = pa.RecordBatch.from_arrays([pa.array(values, type=pa.binary())], names=["b"]) + schema = T.StructType([T.StructField("b", T.BinaryType())]) + _results(service, batch, schema if include_schema else None) + rows = _action(spark.range(1), method) + assert [row.b for row in rows] == values + if values and values[0] is not None: + assert type(rows[0].b) is bytearray + assert service.read_confs == [_BINARY_CONF] + assert service.calls == 0, "binary conversion must not analyze the schema" + + @pytest.mark.parametrize("method", ["collect", "iterator", "prefetch"]) + @pytest.mark.parametrize("include_schema", [False, True]) + def test_schema_free_nested_binary_and_null_first_rows( + self, mock_connect, method, include_schema + ): + pa = pytest.importorskip("pyarrow") + spark, service = mock_connect + service.confs = {_BINARY_CONF: "false"} + batch = pa.RecordBatch.from_arrays([ + pa.array([None, b"x", b""], type=pa.binary()), + pa.array([None, [b"a", None], []], type=pa.list_(pa.binary())), + pa.array([None, {"b": b"s"}, {"b": None}], type=pa.struct([("b", pa.binary())])), + pa.array([None, [("k", b"m")], []], type=pa.map_(pa.string(), pa.binary())), + ], names=["b", "a", "s", "m"]) + schema = T.StructType([ + T.StructField("b", T.BinaryType()), + T.StructField("a", T.ArrayType(T.BinaryType())), + T.StructField("s", T.StructType([T.StructField("b", T.BinaryType())])), + T.StructField("m", T.MapType(T.StringType(), T.BinaryType())), + ]) + _results(service, batch, schema if include_schema else None) + rows = _action(spark.range(3), method) + assert tuple(rows[0]) == (None, None, None, None) + assert all(type(value) is bytearray for value in ( + rows[1].b, rows[1].a[0], rows[1].s.b, rows[1].m["k"], rows[2].b, + )) + assert rows[1].a == [b"a", None] and rows[2].b == b"" + assert rows[2].a == [] and rows[2].s.b is None and rows[2].m == {} + assert service.read_confs == [_BINARY_CONF] and service.calls == 0 + + @pytest.mark.parametrize("setting,expected", [ + ("unset", bytes), ("TrUe", bytes), ("FaLsE", bytearray), + (" true ", bytearray), (None, bytearray), + ]) + def test_result_setting_keeps_reference_default_and_boolean_handling( + self, mock_connect, setting, expected + ): + pa = pytest.importorskip("pyarrow") + spark, service = mock_connect + if setting != "unset": + service.confs = {_BINARY_CONF: setting} + _results(service, pa.RecordBatch.from_arrays([pa.array([b"x"])], names=["b"])) + assert type(spark.range(1).collect()[0].b) is expected + operation = parse_message(parse_message(service.config_requests[0])[3][0]) + pair = parse_message(parse_message(operation[3][0])[1][0]) + assert pair == {1: [_BINARY_CONF.encode()], 2: [b"true"]} + + @pytest.mark.parametrize("prefetch", [False, True]) + def test_iterator_snapshots_on_first_next_before_execution(self, mock_connect, prefetch): + pa = pytest.importorskip("pyarrow") + spark, service = mock_connect + _results(service, pa.RecordBatch.from_arrays([pa.array([b"x", b"y"])], names=["b"])) + service.confs = {_BINARY_CONF: "true"} + service.block_execute = True + iterator = spark.range(2).toLocalIterator(prefetch) + assert service.rpc_events == [] + service.confs = {_BINARY_CONF: "false"} + with ThreadPoolExecutor(max_workers=1) as pool: + pending = pool.submit(next, iterator) + try: + _wait_for(lambda: service.execute_calls == 1) + assert service.rpc_events == ["config", "execute"] + service.confs = {_BINARY_CONF: "true"} + finally: + service.release_response() + assert type(pending.result(timeout=10).b) is bytearray + assert type(next(iterator).b) is bytearray + with pytest.raises(StopIteration): + next(iterator) + assert service.read_confs == [_BINARY_CONF] + + def test_eager_result_reads_setting_after_execution(self, mock_connect): + pa = pytest.importorskip("pyarrow") + spark, service = mock_connect + _results(service, pa.RecordBatch.from_arrays([pa.array([b"x"])], names=["b"])) + service.confs = {_BINARY_CONF: "false"} + service.block_execute = True + frame = spark.range(1) + with ThreadPoolExecutor(max_workers=1) as pool: + pending = pool.submit(frame.collect) + try: + _wait_for(lambda: service.execute_calls == 1) + assert service.rpc_events == ["execute"] + service.confs = {_BINARY_CONF: "true"} + finally: + service.release_response() + assert type(pending.result(timeout=10)[0].b) is bytes + assert service.rpc_events == ["execute", "config"] + + @pytest.mark.parametrize("method", _EAGER_ACTIONS) + def test_eager_configuration_failure_follows_execution(self, mock_connect, method): + pa = pytest.importorskip("pyarrow") + spark, service = mock_connect + _results(service, pa.RecordBatch.from_arrays([pa.array([b"x"])], names=["b"])) + service.config_failure = "binary setting unavailable" + with pytest.raises(PySparkException, match="binary setting unavailable"): + _action(spark.range(1), method) + assert service.rpc_events == ["execute", "config"] + assert service.execute_calls == 1 and len(service.config_requests) == 1 + + def test_eager_execution_error_precedes_configuration(self, mock_connect): + spark, service = mock_connect + service.execute_failure = "execution failed" + service.config_failure = "binary setting unavailable" + with pytest.raises(PySparkException, match="execution failed"): + spark.range(1).collect() + assert service.rpc_events == ["execute"] and service.config_requests == [] + + @pytest.mark.parametrize("prefetch", [False, True]) + def test_iterator_config_failure_is_terminal_without_execution(self, mock_connect, prefetch): + spark, service = mock_connect + service.config_failure = "binary setting unavailable" + iterator = spark.range(2).toLocalIterator(prefetch) + assert service.rpc_events == [] + with pytest.raises(PySparkException, match="binary setting unavailable"): + next(iterator) + service.config_failure = None + with pytest.raises(StopIteration): + next(iterator) + iterator.close() + assert service.rpc_events == ["config"] + assert service.execute_calls == 0 + + @pytest.mark.parametrize("prefetch", [False, True]) + @pytest.mark.parametrize("reason", ["explicit", "converter_failure"]) + def test_iterator_close_releases_execution_before_completion( + self, mock_connect, prefetch, reason + ): + pa = pytest.importorskip("pyarrow") + _, service = mock_connect + if reason == "converter_failure": + spatial = pa.struct([("wkb", pa.binary()), ("srid", pa.int32())]) + batch = pa.RecordBatch.from_arrays([ + pa.array([{"wkb": b"x", "srid": 4326}] * 2, type=spatial), + ], names=["g"]) + schema = T.StructType([T.StructField("g", T.GeometryType(4326))]) + else: + batch = pa.RecordBatch.from_arrays([pa.array([b"x", b"y"])], names=["b"]) + schema = None + _results(service, batch, schema) + service.block_result_complete = True + # Isolate the native close/failure path: retaining the GIL or waiting forever in + # close must fail this bounded test without hanging the pytest interpreter. + child = subprocess.Popen( + [sys.executable, "-c", _ITERATOR_RELEASE_CHILD, service.remote, str(prefetch), reason], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, + ) + try: + _wait_for(lambda: bool(service.execute_requests) or child.poll() is not None, timeout=20) + assert service.execute_requests, "Iterator did not start an execution" + operation_id = parse_message(service.execute_requests[0])[6][0] + + def released(): + # ReleaseExecuteRequest.release_all=5 and operation_id=3. + return any( + 5 in fields and fields.get(3) == [operation_id] + for fields in map(parse_message, service.release_execute_requests) + ) + + _wait_for(lambda: released() or child.poll() is not None) + assert released(), "Closing the iterator did not release its execution" + # ResultComplete cannot have caused this release while r2 is still held. + assert operation_id not in service.result_complete_operations + stdout, stderr = child.communicate("finish\n", timeout=10) + assert child.returncode == 0, stdout + stderr + finally: + service.release_response() + if child.poll() is None: + child.kill() + child.communicate(timeout=10) + assert service.rpc_events == ["config", "execute"] + + @pytest.mark.parametrize("prefetch", [False, True]) + def test_iterator_closed_before_first_next_makes_no_requests(self, mock_connect, prefetch): + spark, service = mock_connect + iterator = spark.range(1).toLocalIterator(prefetch) + iterator.close() + assert list(iterator) == [] + assert service.rpc_events == [] + assert service.calls == 0 + assert service.release_execute_requests == [] + + @pytest.mark.parametrize("iterator", [False, True]) + def test_configuration_wait_releases_gil(self, mock_connect, iterator): + pa = pytest.importorskip("pyarrow") + spark, service = mock_connect + _results(service, pa.RecordBatch.from_arrays([pa.array([b"x"])], names=["b"])) + service.block_config = True + + def release_configuration(): + _wait_for(lambda: bool(service.config_requests)) + service.release_response() + + frame = spark.range(1) + with ThreadPoolExecutor(max_workers=1) as pool: + released = pool.submit(release_configuration) + rows = _action(frame, "iterator" if iterator else "collect") + released.result(timeout=10) + assert type(rows[0].b) is bytes + + @pytest.mark.parametrize("method", ["collect", "iterator", "prefetch"]) + def test_direct_binary_map_keys_keep_reference_unhashable_error(self, mock_connect, method): + pa = pytest.importorskip("pyarrow") + spark, service = mock_connect + service.confs = {_BINARY_CONF: "false"} + batch = pa.RecordBatch.from_arrays([ + pa.array([[(b"key", b"value")]], type=pa.map_(pa.binary(), pa.binary())), + ], names=["m"]) + schema = T.StructType([T.StructField("m", T.MapType(T.BinaryType(), T.BinaryType()))]) + _results(service, batch, schema) + with pytest.raises(TypeError, match="unhashable type: 'bytearray'"): + _action(spark.range(1), method)[0].m diff --git a/python/tests/day_time_interval_cases.py b/python/tests/day_time_interval_cases.py index 11fb422e..c2fe88c7 100644 --- a/python/tests/day_time_interval_cases.py +++ b/python/tests/day_time_interval_cases.py @@ -62,7 +62,8 @@ def _messages(): pool.Add(file) names = ["DataType", "ExecutePlanRequest", "ExecutePlanResponse", "AnalyzePlanRequest", "AnalyzePlanResponse", "ReleaseExecuteRequest", - "ReleaseExecuteResponse", "ReleaseSessionRequest", "ReleaseSessionResponse"] + "ReleaseExecuteResponse", "ReleaseSessionRequest", "ReleaseSessionResponse", + "ConfigRequest", "ConfigResponse"] return {name: message_factory.GetMessageClass( pool.FindMessageTypeByName(f"spark.connect.{name}")) for name in names} @@ -164,6 +165,16 @@ def release_session(raw, context): return self.messages["ReleaseSessionResponse"]( session_id=request.session_id).SerializeToString() + def configure(raw, context): + request = self.messages["ConfigRequest"].FromString(raw) + operation = request.operation.WhichOneof("op_type") + if operation != "get_with_default": + context.abort(grpc.StatusCode.UNIMPLEMENTED, f"unexpected Config operation {operation}") + return self.messages["ConfigResponse"]( + session_id=request.session_id, + pairs=list(request.operation.get_with_default.pairs), + ).SerializeToString() + with ThreadPoolExecutor(max_workers=2) as pool: server = grpc.server(pool) server.add_generic_rpc_handlers([grpc.method_handlers_generic_handler( @@ -172,6 +183,7 @@ def release_session(raw, context): "AnalyzePlan": grpc.unary_unary_rpc_method_handler(analyze), "ReleaseExecute": grpc.unary_unary_rpc_method_handler(release_execute), "ReleaseSession": grpc.unary_unary_rpc_method_handler(release_session), + "Config": grpc.unary_unary_rpc_method_handler(configure), } )]) port = server.add_insecure_port("127.0.0.1:0") diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 60563cd1..9c85da82 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -130,7 +130,7 @@ def release_execute(request, context): ).SerializeToString() def config(request, context): - # The settings that toPandas reads first, as SparkConnectClient.to_pandas does. + # The settings that toPandas reads first, and the defaults that row collection reads. if calls is not None: calls["config"] = calls.get("config", 0) + 1 request = protocol("ConfigRequest").FromString(request) @@ -139,11 +139,23 @@ def config(request, context): "spark.sql.execution.pandas.structHandlingMode": "legacy", "spark.sql.execution.arrow.pyspark.selfDestruct.enabled": "false", } + operation = request.operation.WhichOneof("op_type") + if operation == "get_with_default": + pairs = [] + for pair in request.operation.get_with_default.pairs: + default = pair.value if pair.HasField("value") else None + value = settings.get(pair.key, default) + pairs.append(dict(key=pair.key, **({"value": value} if value is not None else {}))) + else: + if operation != "get": + context.abort(grpc.StatusCode.UNIMPLEMENTED, + f"unexpected Config operation {operation}") + pairs = [dict(key=key, value=settings[key]) for key in request.operation.get.keys] return protocol( "ConfigResponse", session_id=request.session_id, server_side_session_id=server_id, - pairs=[dict(key=key, value=settings[key]) for key in request.operation.get.keys], + pairs=pairs, ).SerializeToString() def release_session(request, context): @@ -227,6 +239,16 @@ def _mixed_batch(): return pa.RecordBatch.from_arrays(arrays, schema=schema) +def test_collect_reads_binary_default_with_arrow_only_schema(protocol): + batch = pa.RecordBatch.from_arrays([pa.array([b"value", None], type=pa.binary())], names=["b"]) + calls = {} + with _dataframe(protocol=protocol, batches=[batch], calls=calls) as (frame, _): + rows = frame.collect() + assert type(rows[0].b) is bytes and rows[0].b == b"value" + assert rows[1].b is None + assert calls == {"execute": 1, "config": 1} + + def test_handoff_preserves_values_metadata_and_chunks(protocol): batch = _mixed_batch() batches = [batch.slice(1, 2), batch.slice(0, 1), batch.slice(3, 1)] @@ -486,18 +508,29 @@ def attribution(request, context): def configure(raw, context): request = protocol("ConfigRequest").FromString(raw) operation = request.operation.WhichOneof("op_type") - keys = list(getattr(request.operation, operation).keys) + defaults = None + if operation == "get_with_default": + defaults = {pair.key: pair.value if pair.HasField("value") else None + for pair in request.operation.get_with_default.pairs} + keys = list(defaults) + else: + if operation != "get": + context.abort(grpc.StatusCode.UNIMPLEMENTED, + f"unexpected Config operation {operation}") + keys = list(request.operation.get.keys) state["config_requests"].append((operation, keys)) if config_error: context.abort(grpc.StatusCode.UNAVAILABLE, "configuration unavailable") # Injected errors abort only the next Config call. if error := state.pop("config_error", None): context.abort(error, "configuration read failed") - # Deliberately reverse responses and omit missing keys to test key mapping. + # Deliberately reverse responses; Get omits missing keys to test key mapping. # A present null key leaves the optional protobuf value unset. - response = protocol("ConfigResponse", session_id=request.session_id, - pairs=[dict(key=k, value=config[k]) for k in reversed(keys) - if k in config]) + if defaults is None: + pairs = [dict(key=k, value=config[k]) for k in reversed(keys) if k in config] + else: + pairs = [dict(key=k, value=config.get(k, defaults[k])) for k in reversed(keys)] + response = protocol("ConfigResponse", session_id=request.session_id, pairs=pairs) state["config_responses"].append(response) return response.SerializeToString() @@ -1034,6 +1067,7 @@ def test_arrow_input_duplicate_chunks_preserve_multiplicity(protocol, batch_byte "spark.sql.session.localRelationChunkSizeBytes", "spark.sql.session.localRelationBatchOfChunksSizeBytes", ] +_BINARY_RESULT_CONFIG = ("get_with_default", ["spark.sql.execution.pyspark.binaryAsBytes"]) @pytest.mark.parametrize("timestamp_type", ["TIMESTAMP_LTZ", "TIMESTAMP_NTZ"]) @@ -1050,6 +1084,8 @@ def test_arrow_input_batched_normalization_configs(protocol, timestamp_type, zon frame = spark.createDataFrame(table) assert state["config_requests"] == [("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] frame.collect() + assert state["config_requests"] == [ + ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS), _BINARY_RESULT_CONFIG] expected_time = (value if timestamp_type == "TIMESTAMP_NTZ" else value.replace(tzinfo=ZoneInfo(zone)).astimezone(timezone.utc)) expected = pa.table({ @@ -1100,7 +1136,7 @@ def test_arrow_input_normalization_config_failure_then_fresh_read(protocol): spark.createDataFrame(table).collect() assert state["config_requests"] == [ ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS), - ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] + ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS), _BINARY_RESULT_CONFIG] assert state["input_tables"][0].to_pydict() == {"x": [1]} @@ -1110,7 +1146,7 @@ def test_arrow_input_normalization_config_transient_retry(protocol): spark.createDataFrame(pa.table({"x": [1]})).collect() assert state["config_requests"] == [ ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS), - ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] + ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS), _BINARY_RESULT_CONFIG] assert state["input_tables"][0].to_pydict() == {"x": [1]} @@ -1125,7 +1161,8 @@ def test_arrow_input_normalization_configs_are_fresh(protocol): state["config"].update(dict(zip(_NORMALIZATION_KEYS, [ "TIMESTAMP_LTZ", "America/Los_Angeles", "false"]))) spark.createDataFrame(table).collect() - assert state["config_requests"] == [("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS)] * 3 + assert state["config_requests"] == [ + ("get", _NORMALIZATION_KEYS + _UPLOAD_KEYS), _BINARY_RESULT_CONFIG] * 3 tables = state["input_tables"] assert [t.column("t")[0].as_py() for t in tables] == [ value.replace(tzinfo=timezone.utc), value, diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 05be95c1..fa5583c3 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -35,6 +35,7 @@ from pyspark.sql import Row, SparkSession, types as T # Collected here so the existing offline CI suite exercises real Arrow results. +from binary_configuration_cases import TestBinaryConfiguration # noqa: F401 from day_time_interval_cases import DayTimeIntervalTests # noqa: F401 from interval_input_cases import IntervalInputConversionTests # noqa: F401 @@ -1637,8 +1638,9 @@ def sent(fn): assert F.udf(plus_one, "int", useArrow=True).evalType == 101 finally: service.confs = {} - # Only udf() without useArrow reads the configuration. - assert service.read_confs == [arrow_conf] * 4 + # Only udf() without useArrow reads its Arrow setting; every collect reads the result setting. + binary_conf = "spark.sql.execution.pyspark.binaryAsBytes" + assert service.read_confs == [arrow_conf, binary_conf, binary_conf] + [arrow_conf, binary_conf] * 3 def test_map_functions_send_the_names_of_their_functions_as_pyspark(mock_connect): @@ -1929,8 +1931,9 @@ def sent(fn): assert sent(F.udtf(Echo, returnType="a int", useArrow=True)) == (False, 301) finally: service.confs = {} - # Only udtf() without useArrow reads the configuration. - assert service.read_confs == [arrow_conf, arrow_conf] + # Only udtf() without useArrow reads its Arrow setting; every collect reads the result setting. + binary_conf = "spark.sql.execution.pyspark.binaryAsBytes" + assert service.read_confs == [arrow_conf] + [binary_conf] * 5 + [arrow_conf] + [binary_conf] * 2 def test_time_values_are_sent_as_pyspark(mock_connect): @@ -9596,7 +9599,6 @@ def test_numpy_scalars_of_other_types_are_unsupported_as_pyspark(mock_connect): (np.uint8(5), "uint8"), (np.uint64(5), "uint64"), (np.float16(1.5), "float16"), - (np.longdouble(1.5), "longdouble"), (np.complex128(1 + 2j), "complex128"), (np.datetime64("2020-01-01"), "datetime64"), (np.timedelta64(1, "D"), "timedelta64"), @@ -9611,6 +9613,90 @@ def test_numpy_scalars_of_other_types_are_unsupported_as_pyspark(mock_connect): assert error.value.getMessageParameters() == {"dtype": "uint8"} +def test_numpy_longdouble_scalar_matches_dtype_and_class_checks_as_pyspark(): + np = pytest.importorskip("numpy") + from pyspark.errors import PySparkTypeError + from pyspark.sql import functions as F + + value = np.longdouble(1.5) + if isinstance(value, (float, np.float32, np.float64)): + pytest.skip("longdouble aliases a supported float scalar class here") + if value.dtype == np.dtype("float64"): + # A dtype mapping to DoubleType does not suffice: PySpark also asserts that its + # scalar is a float, np.float32 or np.float64. Distinct eight-byte longdouble + # classes (such as on macOS arm64 and Windows) fail without a message. + with pytest.raises(AssertionError) as error: + F.lit(value) + assert error.value.args == () + else: + # Linux CI normally has a wider longdouble, which exercises the unsupported + # dtype path; the distinct integer aliases below exercise the class check there. + with pytest.raises(PySparkTypeError) as error: + F.lit(value) + assert error.value.getCondition() == "UNSUPPORTED_DATA_TYPE" + assert error.value.getMessageParameters() == {"data_type": type(value).__name__} + + +def test_numpy_longdouble_array_matches_dtype_and_element_checks_as_pyspark(): + np = pytest.importorskip("numpy") + from pyspark.errors import PySparkTypeError + from pyspark.sql import functions as F + + value = np.array([1.5], dtype=np.longdouble) + if isinstance(value[0], (float, np.float32, np.float64)): + pytest.skip("longdouble aliases a supported float scalar class here") + if value.dtype == np.dtype("float64"): + # A mapped array dtype still builds its literals through each scalar's checks. + with pytest.raises(AssertionError) as error: + F.lit(value) + assert error.value.args == () + else: + with pytest.raises(PySparkTypeError) as error: + F.lit(value) + assert error.value.getCondition() == "UNSUPPORTED_NUMPY_ARRAY_SCALAR" + assert error.value.getMessageParameters() == {"dtype": value.dtype.name} + + +@pytest.mark.parametrize("name", [ + "byte", "short", "intc", "int_", "intp", "long", "longlong", "single", "double", +]) +@pytest.mark.parametrize("as_array", [False, True], ids=["scalar", "array"]) +def test_numpy_numeric_aliases_match_reference_scalar_classes(name, as_array): + np = pytest.importorskip("numpy") + from pyspark.sql import functions as F + + # Read real scalar exports without NumPy's deprecated builtin-alias lookup warnings. + scalar_type = vars(np).get(name) + if scalar_type is None or not issubclass(scalar_type, np.generic): + pytest.skip(f"numpy.{name} is not a NumPy scalar class in this NumPy") + value = scalar_type(1.5 if np.issubdtype(scalar_type, np.floating) else 5) + kind, size = value.dtype.kind, value.dtype.itemsize + if kind == "i": + # These are the classes asserted by LiteralExpression for the inferred width. + accepted_classes = (int,) + tuple( + getattr(np, f"int{width}") for width in (8, 16, 32, 64) if width <= size * 8 + ) + spark_type = {1: "tinyint", 2: "smallint", 4: "int", 8: "bigint"}[size] + rendered = "5" + else: + accepted_classes = (float, np.float32) + ((np.float64,) if size == 8 else ()) + spark_type = {4: "float", 8: "double"}[size] + rendered = "1.5" + literal = np.array([value], dtype=scalar_type) if as_array else value + if not isinstance(value, accepted_classes): + # Distinct longlong scalars have int64 dtype but fail the class assertion on + # Linux/macOS; platforms where it aliases int64 exercise the success case. + with pytest.raises(AssertionError) as error: + F.lit(literal) + assert error.value.args == () + elif as_array: + assert repr(F.lit(literal)).lower() == ( + f"column<'cast(array({rendered}) as array<{spark_type}>)'>" + ) + else: + assert repr(F.lit(literal)) == f"Column<'{rendered}'>" + + # --------------------------------------------------------------------------- 4.3.0 gaps @@ -9988,6 +10074,7 @@ def test_lit_of_a_numpy_array_is_an_array_of_its_type_as_pyspark(): for value, expected in ( (np.array([1, 2]), "cast(array(1, 2) as array)"), (np.array([1.5], dtype=np.float32), "cast(array(1.5) as array)"), + (np.array([1.5], dtype=np.float64), "cast(array(1.5) as array)"), (np.array(["x"]), "cast(array(x) as array)"), ): assert repr(F.lit(value)).lower() == f"column<'{expected}'>"