[SPARK-61099][CONNECT][PYTHON] Preserve binary configuration and iterator cleanup

crates/pyspark-rs/src/dataframe.rs

@@ -4089,43 +4089,33 @@ 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<bool> {
-    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<bool> {
+    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,
     rows: Vec<spark_connect::row::Row>,
     schema: Option<&spark_connect::types::DataType>,
 ) -> PyResult<Vec<PyRow>> {
     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<Option<FieldConverters>>,
-    /// 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<bool>,
 }
 
 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<Option<PyRow>> {
+        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<Option<PyRow>> {
         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
             }
         }
     }

crates/pyspark-rs/src/functions.rs

@@ -333,9 +333,11 @@ fn numpy_ndarray_ndim(obj: &Bound<'_, PyAny>) -> PyResult<Option<usize>> {
 /// 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<Option<LiteralExpression>> {
     let py = obj.py();
     let modules = py.import("sys")?.getattr("modules")?;
@@ -352,6 +354,36 @@ fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult<Option<LiteralExpres
     let dtype = obj.getattr("dtype")?;
     let kind: String = dtype.getattr("kind")?.extract()?;
     let itemsize: usize = dtype.getattr("itemsize")?.extract()?;
+    // Mirror LiteralExpression's accepted scalar classes after dtype inference. Aliases can
+    // share a supported dtype without being instances of its canonical NumPy scalar class.
+    // Check the matching width first; the remaining classes preserve the reference's full set.
+    let scalar_classes: &[&str] = match (kind.as_str(), itemsize) {
+        ("i", 1) => &["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::<pyo3::types::PyInt>(),
+            "f" => obj.is_instance_of::<pyo3::types::PyFloat>(),
+            _ => 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::<i64>()? as i32),

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) {

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),