From 228a511e1dd36202b8e2cff8081ef17f803a2420 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 6 Nov 2026 23:18:54 -0800 Subject: [PATCH 01/20] [SPARK-61102][CONNECT][RUST] Match NumPy longdouble literal class checks --- crates/pyspark-rs/src/functions.rs | 20 +++++++++--- python/tests/test_dropin_offline.py | 47 ++++++++++++++++++++++++++++- 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/crates/pyspark-rs/src/functions.rs b/crates/pyspark-rs/src/functions.rs index 74da2f4f..233d42e2 100644 --- a/crates/pyspark-rs/src/functions.rs +++ b/crates/pyspark-rs/src/functions.rs @@ -333,9 +333,10 @@ 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 distinct longdouble scalar whose dtype maps to double fails the +/// reference's scalar-class assertion. 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")?; @@ -359,7 +360,18 @@ fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult LiteralExpression::Integer(obj.extract::()? as i32), ("i", 8) => LiteralExpression::Long(obj.extract::()?), ("f", 4) => LiteralExpression::Float(obj.extract::()? as f32), - ("f", 8) => LiteralExpression::Double(obj.extract::()?), + ("f", 8) => { + // On some platforms longdouble has an eight-byte dtype equal to float64, but + // remains a distinct scalar class. PySpark infers DoubleType then asserts the + // accepted classes before converting its value to float. + if !obj.is_instance_of::() + && !obj.is_instance(&numpy.getattr("float32")?)? + && !obj.is_instance(&numpy.getattr("float64")?)? + { + return Err(pyo3::exceptions::PyAssertionError::new_err(())); + } + LiteralExpression::Double(obj.extract::()?) + } _ => { let type_name = obj .get_type() diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 05be95c1..f0182860 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -9596,7 +9596,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 +9610,51 @@ 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 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) fail that assertion without a message. + if isinstance(value, (float, np.float32, np.float64)): + assert repr(F.lit(value)) == "Column<'1.5'>" + else: + with pytest.raises(AssertionError) as error: + F.lit(value) + assert error.value.args == () + else: + # Wider longdouble dtypes do not map to a Spark type at all. + 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 value.dtype == np.dtype("float64"): + if isinstance(value[0], (float, np.float32, np.float64)): + assert repr(F.lit(value)).lower() == "column<'cast(array(1.5) as array)'>" + else: + # 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} + + # --------------------------------------------------------------------------- 4.3.0 gaps @@ -9988,6 +10032,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}'>" From 8c216dfd97c74a9e4d5f2eb81b39eb814cea4dce Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 6 Nov 2026 23:29:44 -0800 Subject: [PATCH 02/20] [SPARK-61102][CONNECT][RUST] Apply scalar class checks consistently to numeric aliases --- crates/pyspark-rs/src/functions.rs | 50 +++++++++++++++------- python/tests/test_dropin_offline.py | 66 ++++++++++++++++++++++------- 2 files changed, 86 insertions(+), 30 deletions(-) diff --git a/crates/pyspark-rs/src/functions.rs b/crates/pyspark-rs/src/functions.rs index 233d42e2..14ac379f 100644 --- a/crates/pyspark-rs/src/functions.rs +++ b/crates/pyspark-rs/src/functions.rs @@ -334,9 +334,10 @@ fn numpy_ndarray_ndim(obj: &Bound<'_, PyAny>) -> PyResult> { /// 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, complex, datetime64 or timedelta64) is UNSUPPORTED_DATA_TYPE, -/// named by its Python type. A distinct longdouble scalar whose dtype maps to double fails the -/// reference's scalar-class assertion. NumPy is looked up in `sys.modules`, as a NumPy -/// scalar exists only when NumPy is imported. +/// 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")?; @@ -353,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), @@ -360,18 +391,7 @@ fn numpy_scalar_literal(obj: &Bound<'_, PyAny>) -> PyResult LiteralExpression::Integer(obj.extract::()? as i32), ("i", 8) => LiteralExpression::Long(obj.extract::()?), ("f", 4) => LiteralExpression::Float(obj.extract::()? as f32), - ("f", 8) => { - // On some platforms longdouble has an eight-byte dtype equal to float64, but - // remains a distinct scalar class. PySpark infers DoubleType then asserts the - // accepted classes before converting its value to float. - if !obj.is_instance_of::() - && !obj.is_instance(&numpy.getattr("float32")?)? - && !obj.is_instance(&numpy.getattr("float64")?)? - { - return Err(pyo3::exceptions::PyAssertionError::new_err(())); - } - LiteralExpression::Double(obj.extract::()?) - } + ("f", 8) => LiteralExpression::Double(obj.extract::()?), _ => { let type_name = obj .get_type() diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index f0182860..749d322c 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -9616,18 +9616,18 @@ def test_numpy_longdouble_scalar_matches_dtype_and_class_checks_as_pyspark(): 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) fail that assertion without a message. - if isinstance(value, (float, np.float32, np.float64)): - assert repr(F.lit(value)) == "Column<'1.5'>" - else: - with pytest.raises(AssertionError) as error: - F.lit(value) - assert error.value.args == () + # 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: - # Wider longdouble dtypes do not map to a Spark type at all. + # 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" @@ -9640,14 +9640,13 @@ def test_numpy_longdouble_array_matches_dtype_and_element_checks_as_pyspark(): 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"): - if isinstance(value[0], (float, np.float32, np.float64)): - assert repr(F.lit(value)).lower() == "column<'cast(array(1.5) as array)'>" - else: - # 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 == () + # 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) @@ -9655,6 +9654,43 @@ def test_numpy_longdouble_array_matches_dtype_and_element_checks_as_pyspark(): assert error.value.getMessageParameters() == {"dtype": value.dtype.name} +@pytest.mark.parametrize("name", [ + "byte", "short", "intc", "int_", "intp", "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 + + scalar_type = getattr(np, name) + 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 From 07b5866057d37db1b8d8f871d9102e1ceaa4867b Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 6 Nov 2026 23:37:57 -0800 Subject: [PATCH 03/20] Cover optional NumPy long scalar aliases --- python/tests/test_dropin_offline.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 749d322c..7e2c0aba 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -9655,14 +9655,17 @@ def test_numpy_longdouble_array_matches_dtype_and_element_checks_as_pyspark(): @pytest.mark.parametrize("name", [ - "byte", "short", "intc", "int_", "intp", "longlong", "single", "double", + "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 - scalar_type = getattr(np, name) + # 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": From 61b93ba243fb2914e258e2096f5f2dedc4f08e4f Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 22:36:45 -0800 Subject: [PATCH 04/20] [SPARK-61099][CONNECT][PYTHON] Preserve binary result configuration and iterator failure semantics --- crates/pyspark-rs/src/dataframe.rs | 72 +++--- python/tests/_mock_connect_server.py | 31 ++- python/tests/binary_configuration_cases.py | 278 +++++++++++++++++++++ python/tests/day_time_interval_cases.py | 12 +- python/tests/test_arrow_handoff.py | 47 +++- python/tests/test_dropin_offline.py | 11 +- 6 files changed, 402 insertions(+), 49 deletions(-) create mode 100644 python/tests/binary_configuration_cases.py diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 88141df3..398a6104 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -4089,28 +4089,17 @@ 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"))) } @@ -4123,9 +4112,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 +4135,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 +4147,7 @@ impl PyLocalRowIterator { iterator: None, finished: false, converters: None, - binary_as_bytes: true, + binary_as_bytes: None, } } @@ -4196,18 +4185,35 @@ impl PyLocalRowIterator { if self.finished { return Ok(None); } + let binary_as_bytes = match self.binary_as_bytes { + Some(value) => value, + None => match read_binary_as_bytes(py, &self.dataframe) { + Ok(value) => { + self.binary_as_bytes = Some(value); + value + } + Err(error) => { + self.close(py); + return Err(error); + } + }, + }; // 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; + self.close(py); 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 udt_field_converters(py, Some(&schema)) { + Ok(converters) => self.converters = Some(converters), + Err(error) => { + self.close(py); + return Err(error); + } + } } } match next { @@ -4216,16 +4222,16 @@ impl PyLocalRowIterator { Ok(Some(PyRow::with_converters( row, converters, - self.binary_as_bytes, + binary_as_bytes, ))) } // As PySpark raises the error of the execution, which ends the iteration. Some(Err(error)) => { - self.finished = true; + self.close(py); Err(crate::errors::spark_error_to_py_exception(error)) } None => { - self.finished = true; + self.close(py); Ok(None) } } diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index 0d4beee0..1011d14e 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -269,6 +269,7 @@ def reset(self): self.block_execute = False self.analyze_requests = [] self.execute_requests = [] + self.release_execute_requests = [] self.execute_payload = b"" self.execute_failure = None self.execute_error = None @@ -280,6 +281,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 +336,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,11 +368,16 @@ 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 self.execute_requests.append(request) + self.rpc_events.append("execute") if self.block_execute and not self.response_gate.wait(5): context.abort( grpc.StatusCode.DEADLINE_EXCEEDED, @@ -408,6 +419,10 @@ 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 +444,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: @@ -457,7 +481,10 @@ def config(self, request, context): 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))) + 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 +512,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..15480faa --- /dev/null +++ b/python/tests/binary_configuration_cases.py @@ -0,0 +1,278 @@ +# 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 time + +import pytest + +from _mock_connect_server import data_type_proto, message_field, parse_message, varint_field +from pyspark.sql import types as T + + +_BINARY_CONF = "spark.sql.execution.pyspark.binaryAsBytes" +_ACTIONS = ["collect", "first", "head", "head_n", "take", "tail", "iterator", "prefetch"] + + +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): + deadline = time.monotonic() + 5 + 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", _ACTIONS[:6]) + 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(Exception, 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(Exception, 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(Exception, 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]) + def test_iterator_converter_initialization_failure_closes_execution( + self, mock_connect, monkeypatch, prefetch + ): + pa = pytest.importorskip("pyarrow") + spark, service = mock_connect + 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"]) + _results(service, batch, T.StructType([T.StructField("g", T.GeometryType(4326))])) + original = T._udt_field_converters + + def fail(schema): + raise LookupError("converter initialization failed") + + monkeypatch.setattr(T, "_udt_field_converters", fail) + iterator = spark.range(2).toLocalIterator(prefetch) + try: + with pytest.raises(LookupError, match="converter initialization failed"): + next(iterator) + monkeypatch.setattr(T, "_udt_field_converters", original) + with pytest.raises(StopIteration): + next(iterator) + operation_id = parse_message(service.execute_requests[0])[6] + # ReleaseExecuteRequest.release_all=5; a closed iterator keeps no server execution. + _wait_for(lambda: any( + 5 in fields and fields.get(3) == operation_id + for fields in map(parse_message, service.release_execute_requests) + )) + finally: + iterator.close() + 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 == [] + + @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..f8f9717e 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,14 @@ 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) + self.assertEqual(request.operation.WhichOneof("op_type"), "get_with_default") + 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 +181,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..4215a0f1 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,18 @@ 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 = [dict(key=pair.key, value=settings.get(pair.key, pair.value)) + for pair in request.operation.get_with_default.pairs] + else: + assert operation == "get" + 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 +234,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 +503,26 @@ 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: + assert operation == "get" + 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]) + pairs=[dict(key=k, value=config.get(k, defaults[k]) if defaults is not None + else config[k]) for k in reversed(keys) + if k in config or defaults is not None]) state["config_responses"].append(response) return response.SerializeToString() @@ -1034,6 +1059,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 +1076,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 +1128,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 +1138,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 +1153,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 7e2c0aba..ce987b5f 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -37,6 +37,7 @@ # Collected here so the existing offline CI suite exercises real Arrow results. from day_time_interval_cases import DayTimeIntervalTests # noqa: F401 from interval_input_cases import IntervalInputConversionTests # noqa: F401 +from binary_configuration_cases import TestBinaryConfiguration # noqa: F401 def test_create_dataframe_input_names_offline(): @@ -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): From 9e7e7d36b3fc88259356b8ba9b75f35b9fb04dd7 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 22:50:35 -0800 Subject: [PATCH 05/20] [SPARK-61099][CONNECT][PYTHON] Release incomplete prefetched executions and tighten lifecycle coverage --- crates/pyspark-rs/src/dataframe.rs | 89 +++++++-------- crates/spark-connect/src/dataframe.rs | 8 +- python/tests/_mock_connect_server.py | 19 ++++ python/tests/binary_configuration_cases.py | 122 +++++++++++++++------ python/tests/day_time_interval_cases.py | 4 +- python/tests/test_arrow_handoff.py | 9 +- python/tests/test_dropin_offline.py | 2 +- 7 files changed, 164 insertions(+), 89 deletions(-) diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 398a6104..e02722aa 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -4104,7 +4104,8 @@ fn read_binary_as_bytes(py: Python<'_>, dataframe: &DataFrame) -> PyResult } /// 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, @@ -4161,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] @@ -4180,59 +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); } - let binary_as_bytes = match self.binary_as_bytes { - Some(value) => value, - None => match read_binary_as_bytes(py, &self.dataframe) { - Ok(value) => { - self.binary_as_bytes = Some(value); - value - } - Err(error) => { - self.close(py); - return Err(error); - } - }, - }; - // 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.close(py); - return Err(error); - } - match udt_field_converters(py, Some(&schema)) { - Ok(converters) => self.converters = Some(converters), - Err(error) => { - self.close(py); - return Err(error); - } - } - } - } - match next { - Some(Ok(row)) => { - let converters = self.converters.clone().flatten(); - Ok(Some(PyRow::with_converters( - row, - converters, - binary_as_bytes, - ))) - } - // As PySpark raises the error of the execution, which ends the iteration. - Some(Err(error)) => { - self.close(py); - Err(crate::errors::spark_error_to_py_exception(error)) - } - None => { + match self.next_python_row(py) { + Ok(Some(row)) => Ok(Some(row)), + result => { self.close(py); - Ok(None) + result } } } diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index bafdebba..ff79f0c0 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -212,7 +212,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 1011d14e..3801da66 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -254,6 +254,11 @@ 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_waiting = False + result_complete_sent = False # Each AnalyzePlanRequest, recorded. analyze_requests = [] @@ -267,6 +272,9 @@ def reset(self): self.block_response = False self.response_gate.clear() self.block_execute = False + self.block_result_complete = False + self.result_complete_waiting = False + self.result_complete_sent = False self.analyze_requests = [] self.execute_requests = [] self.release_execute_requests = [] @@ -413,6 +421,17 @@ 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 self.block_result_complete: + self.result_complete_waiting = True + try: + if not self.response_gate.wait(30): + context.abort( + grpc.StatusCode.DEADLINE_EXCEEDED, + "ResultComplete was not released by the test", + ) + finally: + self.result_complete_waiting = False + self.result_complete_sent = True yield header + message_field(13, "r2") + message_field(14, b"") def warnings_fields(self): diff --git a/python/tests/binary_configuration_cases.py b/python/tests/binary_configuration_cases.py index 15480faa..6fb3820c 100644 --- a/python/tests/binary_configuration_cases.py +++ b/python/tests/binary_configuration_cases.py @@ -17,16 +17,54 @@ """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" -_ACTIONS = ["collect", "first", "head", "head_n", "take", "tail", "iterator", "prefetch"] +_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): @@ -53,8 +91,8 @@ def _action(frame, method): return list(frame.toLocalIterator(method == "prefetch")) -def _wait_for(predicate): - deadline = time.monotonic() + 5 +def _wait_for(predicate, timeout=5): + deadline = time.monotonic() + timeout while not predicate(): if time.monotonic() >= deadline: raise AssertionError("mock RPC did not arrive") @@ -170,13 +208,13 @@ def test_eager_result_reads_setting_after_execution(self, mock_connect): assert type(pending.result(timeout=10)[0].b) is bytes assert service.rpc_events == ["execute", "config"] - @pytest.mark.parametrize("method", _ACTIONS[:6]) + @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(Exception, match="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 @@ -185,7 +223,7 @@ 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(Exception, match="execution failed"): + with pytest.raises(PySparkException, match="execution failed"): spark.range(1).collect() assert service.rpc_events == ["execute"] and service.config_requests == [] @@ -195,7 +233,7 @@ def test_iterator_config_failure_is_terminal_without_execution(self, mock_connec service.config_failure = "binary setting unavailable" iterator = spark.range(2).toLocalIterator(prefetch) assert service.rpc_events == [] - with pytest.raises(Exception, match="binary setting unavailable"): + with pytest.raises(PySparkException, match="binary setting unavailable"): next(iterator) service.config_failure = None with pytest.raises(StopIteration): @@ -205,37 +243,53 @@ def test_iterator_config_failure_is_terminal_without_execution(self, mock_connec assert service.execute_calls == 0 @pytest.mark.parametrize("prefetch", [False, True]) - def test_iterator_converter_initialization_failure_closes_execution( - self, mock_connect, monkeypatch, prefetch + @pytest.mark.parametrize("reason", ["explicit", "converter_failure"]) + def test_iterator_close_releases_execution_before_completion( + self, mock_connect, prefetch, reason ): pa = pytest.importorskip("pyarrow") - spark, service = mock_connect - 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"]) - _results(service, batch, T.StructType([T.StructField("g", T.GeometryType(4326))])) - original = T._udt_field_converters - - def fail(schema): - raise LookupError("converter initialization failed") - - monkeypatch.setattr(T, "_udt_field_converters", fail) - iterator = spark.range(2).toLocalIterator(prefetch) + _, 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: - with pytest.raises(LookupError, match="converter initialization failed"): - next(iterator) - monkeypatch.setattr(T, "_udt_field_converters", original) - with pytest.raises(StopIteration): - next(iterator) + _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] - # ReleaseExecuteRequest.release_all=5; a closed iterator keeps no server execution. - _wait_for(lambda: any( - 5 in fields and fields.get(3) == operation_id - for fields in map(parse_message, service.release_execute_requests) - )) + + 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 service.block_result_complete + assert not service.result_complete_sent + stdout, stderr = child.communicate("finish\n", timeout=10) + assert child.returncode == 0, stdout + stderr finally: - iterator.close() + 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]) @@ -245,6 +299,8 @@ def test_iterator_closed_before_first_next_makes_no_requests(self, mock_connect, 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): diff --git a/python/tests/day_time_interval_cases.py b/python/tests/day_time_interval_cases.py index f8f9717e..c2fe88c7 100644 --- a/python/tests/day_time_interval_cases.py +++ b/python/tests/day_time_interval_cases.py @@ -167,7 +167,9 @@ def release_session(raw, context): def configure(raw, context): request = self.messages["ConfigRequest"].FromString(raw) - self.assertEqual(request.operation.WhichOneof("op_type"), "get_with_default") + 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), diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 4215a0f1..ed9d6407 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -519,10 +519,11 @@ def configure(raw, context): context.abort(error, "configuration read failed") # 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.get(k, defaults[k]) if defaults is not None - else config[k]) for k in reversed(keys) - if k in config or defaults is not None]) + 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() diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index ce987b5f..fa5583c3 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -35,9 +35,9 @@ 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 -from binary_configuration_cases import TestBinaryConfiguration # noqa: F401 def test_create_dataframe_input_names_offline(): From 1f5b90927ef26d2cfab35660ec74c2d71422831b Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 22:59:40 -0800 Subject: [PATCH 06/20] [SPARK-61099][CONNECT][PYTHON] Isolate incomplete execution gates and completion records per operation --- crates/spark-connect/src/dataframe.rs | 9 +++--- python/tests/_mock_connect_server.py | 32 ++++++++++------------ python/tests/binary_configuration_cases.py | 7 ++--- 3 files changed, 23 insertions(+), 25 deletions(-) diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index ff79f0c0..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, diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index 3801da66..d3ca2f4d 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -257,8 +257,7 @@ class Service: # 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_waiting = False - result_complete_sent = False + result_complete_operations = [] # Each AnalyzePlanRequest, recorded. analyze_requests = [] @@ -270,11 +269,11 @@ 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_waiting = False - self.result_complete_sent = False + self.result_complete_operations = [] self.analyze_requests = [] self.execute_requests = [] self.release_execute_requests = [] @@ -383,10 +382,14 @@ def analyze(self, request, context): set_confs = [] def execute(self, request, context): + # 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 self.execute_calls += 1 self.execute_requests.append(request) self.rpc_events.append("execute") - if self.block_execute and not self.response_gate.wait(5): + 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", @@ -421,17 +424,12 @@ 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 self.block_result_complete: - self.result_complete_waiting = True - try: - if not self.response_gate.wait(30): - context.abort( - grpc.StatusCode.DEADLINE_EXCEEDED, - "ResultComplete was not released by the test", - ) - finally: - self.result_complete_waiting = False - self.result_complete_sent = True + 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): diff --git a/python/tests/binary_configuration_cases.py b/python/tests/binary_configuration_cases.py index 6fb3820c..372b45b5 100644 --- a/python/tests/binary_configuration_cases.py +++ b/python/tests/binary_configuration_cases.py @@ -269,20 +269,19 @@ def test_iterator_close_releases_execution_before_completion( 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] + 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 + 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 service.block_result_complete - assert not service.result_complete_sent + assert operation_id not in service.result_complete_operations stdout, stderr = child.communicate("finish\n", timeout=10) assert child.returncode == 0, stdout + stderr finally: From 843df17b6e36a71725dab48cf0512418c97d7162 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 23:52:48 -0800 Subject: [PATCH 07/20] Preserve unset defaults in the collection test service --- python/tests/test_arrow_handoff.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index ed9d6407..604d5d77 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -141,8 +141,11 @@ def config(request, context): } operation = request.operation.WhichOneof("op_type") if operation == "get_with_default": - pairs = [dict(key=pair.key, value=settings.get(pair.key, pair.value)) - for pair in request.operation.get_with_default.pairs] + 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: assert operation == "get" pairs = [dict(key=key, value=settings[key]) for key in request.operation.get.keys] From fc5c61331111509225a845a1e761b48fedd54721 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 23:59:47 -0800 Subject: [PATCH 08/20] Keep configuration mock defaults and errors consistent --- python/tests/_mock_connect_server.py | 7 +++++-- python/tests/test_arrow_handoff.py | 8 ++++++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index d3ca2f4d..de1b606a 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -491,13 +491,16 @@ 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()) + 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) diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 604d5d77..9c85da82 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -147,7 +147,9 @@ def config(request, context): value = settings.get(pair.key, default) pairs.append(dict(key=pair.key, **({"value": value} if value is not None else {}))) else: - assert operation == "get" + 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", @@ -512,7 +514,9 @@ def configure(raw, context): for pair in request.operation.get_with_default.pairs} keys = list(defaults) else: - assert operation == "get" + 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: From 12bd26f9c3cc96524efa92025ea2da6291442dc9 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sun, 8 Nov 2026 00:07:46 -0800 Subject: [PATCH 09/20] [SPARK-61099][CONNECT][TESTS] Publish execution bookkeeping before the readiness counter --- python/tests/_mock_connect_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index de1b606a..a150c754 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -386,9 +386,9 @@ def execute(self, request, context): response_gate = self.response_gate block_result_complete = self.block_result_complete result_complete_operations = self.result_complete_operations - self.execute_calls += 1 self.execute_requests.append(request) 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, From 8ea89ffe5ced04da601e31a538687601c0910671 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sun, 8 Nov 2026 00:13:53 -0800 Subject: [PATCH 10/20] [SPARK-61099][CONNECT][TESTS] Document cross-thread execution readiness ordering --- python/tests/_mock_connect_server.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index a150c754..38f49653 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -386,6 +386,7 @@ def execute(self, request, context): 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) self.rpc_events.append("execute") self.execute_calls += 1 From 3a8683a70c1266749d2ef6bc6903443623e9413d Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 6 Nov 2026 23:27:36 -0800 Subject: [PATCH 11/20] [SPARK-61100][CONNECT][RUST] Return tuple-backed Python Rows --- .github/workflows/lint.yml | 6 +- AGENTS.md | 16 +- crates/pyspark-rs/src/dataframe.rs | 95 +++-- crates/pyspark-rs/src/infer.rs | 22 +- crates/pyspark-rs/src/lib.rs | 3 +- crates/pyspark-rs/src/row.py | 140 +++++++ crates/pyspark-rs/src/row.rs | 601 ++++++---------------------- crates/pyspark-rs/src/session.rs | 30 +- crates/spark-connect/src/row.rs | 11 + pyproject.toml | 10 +- python/.coveragerc | 2 + python/pyspark/sql/types.py | 12 +- python/tests/test_arrow_handoff.py | 92 +++++ python/tests/test_dropin_offline.py | 299 +++++++++++++- 14 files changed, 780 insertions(+), 559 deletions(-) create mode 100644 crates/pyspark-rs/src/row.py diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 1f226e29..c1473891 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -66,7 +66,7 @@ jobs: run: python3 -m unittest discover -s scripts/tests python-ruff: - name: ruff (our scripts) + name: ruff (our scripts and embedded Row) runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 @@ -76,6 +76,6 @@ jobs: - name: Install ruff run: python -m pip install "ruff==0.14.0" - name: ruff check - blocking - run: ruff check scripts/ + run: ruff check scripts/ crates/pyspark-rs/src/row.py - name: ruff format (check) - blocking - run: ruff format --check scripts/ + run: ruff format --check scripts/ crates/pyspark-rs/src/row.py diff --git a/AGENTS.md b/AGENTS.md index 4810ae4d..612004c3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -243,7 +243,9 @@ Vendor the ones we don't have a Rust implementation for. (the gRPC channel connects lazily, so proto/plan construction needs no server) and server-gated `tests/e2e_*.rs` (they check `std::env::var("SPARK_REMOTE")`). - **Python drop-in**: `python/.coveragerc` scopes to the hand-written shim - (omits vendored/generated/Rust-only). Run: + (omits vendored/generated/Rust-only). The embedded `crates/pyspark-rs/src/row.py` + executes under a synthetic filename outside this coverage scope; its behavior is + exercised by the Row tests and its source is included in Ruff. Run: ``` cd python && PYTHONPATH=$(pwd) python -m coverage run --rcfile=.coveragerc \ -m pytest tests/test_dropin_offline.py && python -m coverage report --rcfile=.coveragerc @@ -414,9 +416,15 @@ For every method matching a pyspark name: proto carrying the pickled object: `registerDataSource`, `GroupedData.transformWithState[InPandas]`, `createDataFrame(..., udf)`, etc. General rule: proto-building → Rust; cloudpickled Python code that runs server-side → vendored Python. -- **Row materialization:** PyO3 cannot subclass `tuple`, so our `Row` is not a tuple subclass. - Some pandas-on-Spark code assumes a tuple-`Row` with mutable `__fields__`. Full fidelity - likely requires adopting the upstream Python tuple-`Row` at collect boundaries. +- **Row materialization:** `crates/pyspark-rs/src/row.py` embeds PySpark's Python `Row(tuple)` + class, exported as `pyspark._pyspark.Row` and re-exported by `pyspark.sql.types`, + `pyspark.sql`, and `pyspark`. Embedding avoids importing the partially initialized + `pyspark.sql.types` shim while the extension loads. Rust eagerly builds tuple values + through `pyspark.sql.types._create_row`; collected rows share a mutable `__fields__` + list across the result, and iterator rows share top-level names within each Arrow + record batch. Nested names are shared by schema path across the execution, including + iterator batches, while distinct sibling paths keep separate lists. + Nested local-input Rows follow the tuple inference paths; named Rows expose `__fields__`. - **Streaming listeners and progress holders:** Keep progress holders (`StreamingQueryProgress`, `StateOperatorProgress`, etc.) as faithful pure-Python implementations parsed from server JSON, not Rust-backed. The listener bus (subscribe/dispatch) is genuinely Rust-backed; diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index e02722aa..8a0d80fb 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -9,7 +9,7 @@ use spark_connect::udf::{eval_type, CommonInlineUserDefinedFunctionExpression, P use crate::column::PyColumn; use crate::errors::ResultExt; use crate::group::PyGroupedData; -use crate::row::{FieldConverters, PyRow}; +use crate::row::{FieldConverters, NestedFieldNames}; use crate::streaming::PyDataStreamWriter; use crate::types::PyDataType; @@ -2048,12 +2048,12 @@ impl PyDataFrame { } /// Collect all rows into memory. - fn collect(&self, py: Python<'_>) -> PyResult> { + fn collect(&self, py: Python<'_>) -> PyResult>> { // Release the GIL across the blocking RPC so other Python threads run. let (rows, schema) = py .detach(|| self.dataframe.collect_with_schema()) .to_pyerr()?; - rows_to_py(py, &self.dataframe, rows, schema.as_ref()) + rows_to_py(py, &self.dataframe, &rows, schema.as_ref()) } /// Return an iterator over rows without materializing the entire result in memory. @@ -2257,12 +2257,11 @@ impl PyDataFrame { } /// Get the first row. - fn first(&self, py: Python<'_>) -> PyResult> { + fn first(&self, py: Python<'_>) -> PyResult>> { let (row, schema) = py .detach(|| self.dataframe.first_with_schema()) .to_pyerr()?; - let rows = row.into_iter().collect(); - Ok(rows_to_py(py, &self.dataframe, rows, schema.as_ref())?.pop()) + Ok(rows_to_py(py, &self.dataframe, row.as_slice(), schema.as_ref())?.pop()) } /// Alias for first. @@ -2277,17 +2276,17 @@ impl PyDataFrame { let (rows, schema) = py .detach(|| self.dataframe.limit(k).collect_with_schema()) .to_pyerr()?; - rows_to_py(py, &self.dataframe, rows, schema.as_ref())?.into_py_any(py) + rows_to_py(py, &self.dataframe, &rows, schema.as_ref())?.into_py_any(py) } } } /// Get the first n rows. - fn take(&self, py: Python<'_>, num: usize) -> PyResult> { + fn take(&self, py: Python<'_>, num: usize) -> PyResult>> { let (rows, schema) = py .detach(|| self.dataframe.take_with_schema(num)) .to_pyerr()?; - rows_to_py(py, &self.dataframe, rows, schema.as_ref()) + rows_to_py(py, &self.dataframe, &rows, schema.as_ref()) } /// Column names (a property, mirroring `DataFrame.columns`). @@ -2297,11 +2296,11 @@ impl PyDataFrame { } /// Get the last `num` rows as a list of Rows (eager, like pyspark `tail`). - fn tail(&self, py: Python<'_>, num: i32) -> PyResult> { + fn tail(&self, py: Python<'_>, num: i32) -> PyResult>> { let (rows, schema) = py .detach(|| self.dataframe.tail(num).collect_with_schema()) .to_pyerr()?; - rows_to_py(py, &self.dataframe, rows, schema.as_ref()) + rows_to_py(py, &self.dataframe, &rows, schema.as_ref()) } /// Sample a fraction of rows, as PySpark's `DataFrame.sample(withReplacement=None, @@ -4103,23 +4102,36 @@ fn read_binary_as_bytes(py: Python<'_>, dataframe: &DataFrame) -> PyResult 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. Reads the binary result setting after execution -/// and applies it to binary values, including when only an Arrow schema is available. +/// The collected `rows` of a result of `schema`, with configured binary values and +/// user-defined values materialized before constructing the tuples returned to Python. +/// Reads the binary setting after execution, including with only an Arrow schema. fn rows_to_py( py: Python<'_>, dataframe: &DataFrame, - rows: Vec, + rows: &[spark_connect::row::Row], schema: Option<&spark_connect::types::DataType>, -) -> PyResult> { +) -> 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)?; - Ok(rows - .into_iter() - .map(|row| PyRow::with_converters(row, converters.clone(), binary_as_bytes)) - .collect()) + let Some(first) = rows.first() else { + return Ok(Vec::new()); + }; + let fields = PyList::new(py, first.fields())?; + let mut nested_fields = NestedFieldNames::default(); + rows.iter() + .map(|row| { + crate::row::row_values_to_py( + py, + &fields, + row.values(), + converters.as_ref(), + binary_as_bytes, + &mut nested_fields, + ) + }) + .collect() } /// Python wrapper for LocalRowIterator yielding Row objects lazily. @@ -4138,6 +4150,10 @@ pub struct PyLocalRowIterator { converters: Option>, /// Snapshotted before the first fetch, as the body of PySpark's generator does. binary_as_bytes: Option, + /// Top-level names shared only by rows of the current Arrow record batch. + fields: Option<(std::sync::Arc<[String]>, Py)>, + /// Nested names belong to schema paths and stay shared across iterator batches. + nested_fields: NestedFieldNames, } impl PyLocalRowIterator { @@ -4149,6 +4165,8 @@ impl PyLocalRowIterator { finished: false, converters: None, binary_as_bytes: None, + fields: None, + nested_fields: NestedFieldNames::default(), } } @@ -4164,7 +4182,7 @@ impl PyLocalRowIterator { } /// Prepare and fetch one Python row; the public wrapper closes on exhaustion or error. - fn next_python_row(&mut self, py: Python<'_>) -> PyResult> { + fn next_python_row(&mut self, py: Python<'_>) -> PyResult>> { let binary_as_bytes = match self.binary_as_bytes { Some(value) => value, None => { @@ -4184,11 +4202,34 @@ impl PyLocalRowIterator { match next { Some(Ok(row)) => { let converters = self.converters.clone().flatten(); - Ok(Some(PyRow::with_converters( - row, - converters, + let fields = match row.shared_fields() { + Some(native_fields) => match self.fields.as_ref() { + Some((cached, fields)) if std::sync::Arc::ptr_eq(cached, native_fields) => { + fields.bind(py).clone() + } + _ => { + let fields = PyList::new(py, row.fields())?; + self.fields = Some(( + std::sync::Arc::clone(native_fields), + fields.clone().unbind(), + )); + fields + } + }, + None => { + self.fields = None; + PyList::new(py, row.fields())? + } + }; + crate::row::row_values_to_py( + py, + &fields, + row.values(), + converters.as_ref(), binary_as_bytes, - ))) + &mut self.nested_fields, + ) + .map(Some) } Some(Err(error)) => Err(crate::errors::spark_error_to_py_exception(error)), None => Ok(None), @@ -4208,6 +4249,8 @@ impl PyLocalRowIterator { fn close(&mut self, py: Python<'_>) { let iterator = self.iterator.take(); self.finished = true; + self.fields = None; + self.nested_fields.clear(); py.detach(move || drop(iterator)); } @@ -4215,7 +4258,7 @@ impl PyLocalRowIterator { /// /// 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> { + fn __next__(&mut self, py: Python<'_>) -> PyResult>> { if self.finished { return Ok(None); } diff --git a/crates/pyspark-rs/src/infer.rs b/crates/pyspark-rs/src/infer.rs index 0234547d..b106b7df 100644 --- a/crates/pyspark-rs/src/infer.rs +++ b/crates/pyspark-rs/src/infer.rs @@ -23,7 +23,6 @@ use spark_connect::types::{DataType, GeographyType, GeometryType, StructField}; use crate::column::raise_pyspark; use crate::errors::ResultExt; -use crate::row::PyRow; use crate::values::{PyGeography, PyGeometry, PyVariantVal}; pub(crate) static DECIMAL: PyOnceLock> = PyOnceLock::new(); @@ -533,18 +532,6 @@ fn row_items<'py>( if row.is_instance_of::() { return sorted_items(row); } - if let Ok(pyrow) = row.cast::() { - let pyrow = pyrow.borrow(); - let (fields, values) = pyrow.parts(py)?; - return Ok(match fields { - Some(fields) => fields - .iter() - .map(|f| PyString::new(py, f).into_any()) - .zip(values.iter()) - .collect(), - None => positional(py, values.iter().collect(), names), - }); - } if row.is_instance_of::() || row.is_instance_of::() { let fields = if row.hasattr(intern!(py, "__fields__"))? { Some(row.getattr(intern!(py, "__fields__"))?) @@ -1533,10 +1520,7 @@ fn values_by_position<'py>( /// and another object by the names of its `__dict__`. fn convert_struct(obj: &Bound<'_, PyAny>, fields: &[StructField]) -> PyResult { let py = obj.py(); - let values = if let Ok(row) = obj.cast::() { - let (_, values) = row.borrow().parts(py)?; - values_by_position(py, values.iter().collect(), fields)? - } else if let Ok(tuple) = obj.cast::() { + let values = if let Ok(tuple) = obj.cast::() { values_by_position(py, tuple.iter().collect(), fields)? } else if obj.is_instance_of::() { values_by_name(obj, fields)? @@ -1592,10 +1576,6 @@ pub(crate) fn row_values<'py>( if item.is_none() { return Ok(vec![py.None().into_bound(py); fields.len()]); } - if let Ok(row) = item.cast::() { - let (_, values) = row.borrow().parts(py)?; - return values_by_position(py, values.iter().collect(), fields); - } if let Ok(tuple) = item.cast::() { return values_by_position(py, tuple.iter().collect(), fields); } diff --git a/crates/pyspark-rs/src/lib.rs b/crates/pyspark-rs/src/lib.rs index 840843af..72c22182 100644 --- a/crates/pyspark-rs/src/lib.rs +++ b/crates/pyspark-rs/src/lib.rs @@ -74,7 +74,6 @@ use readwriter::PyDataFrameReader; use resource::{ PyExecutorResourceRequests, PyResourceProfile, PyResourceProfileBuilder, PyTaskResourceRequests, }; -use row::PyRow; use session::{PyHook, PySparkSession, PySparkSessionBuilder}; use stat::PyStatFunctions; use streaming::{ @@ -109,7 +108,7 @@ fn _pyspark(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; + row::register(m)?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/crates/pyspark-rs/src/row.py b/crates/pyspark-rs/src/row.py new file mode 100644 index 00000000..6abb3221 --- /dev/null +++ b/crates/pyspark-rs/src/row.py @@ -0,0 +1,140 @@ +# +# 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. +# + +"""Tuple-backed Row semantics from PySpark's ``pyspark.sql.types.Row``. + +Mirrors ``class Row`` in apache/spark v4.3.0, ``python/pyspark/sql/types.py``. +Embedded by the extension so ``pyspark._pyspark.Row`` and the public imports +refer to the same class without importing the partially initialized types shim. +""" + + +class Row(tuple): + """A row with optional field names, also usable as a factory of named rows.""" + + __module__ = "pyspark.sql.types" + + def __new__(cls, *args, **kwargs): + if args and kwargs: + from pyspark.errors import PySparkValueError + + raise PySparkValueError( + errorClass="CANNOT_SET_TOGETHER", + messageParameters={"arg_list": "args and kwargs"}, + ) + if kwargs: + row = tuple.__new__(cls, list(kwargs.values())) + row.__fields__ = list(kwargs.keys()) + return row + return tuple.__new__(cls, args) + + def asDict(self, recursive=False): + """Return field names and values as a dict, optionally converting nested rows.""" + if not hasattr(self, "__fields__"): + from pyspark.errors import PySparkTypeError + + raise PySparkTypeError( + errorClass="CANNOT_CONVERT_TYPE", + messageParameters={"from_type": "Row", "to_type": "dict"}, + ) + if recursive: + + def conv(obj): + if isinstance(obj, Row): + return obj.asDict(True) + if isinstance(obj, list): + return [conv(value) for value in obj] + if isinstance(obj, dict): + return {key: conv(value) for key, value in obj.items()} + return obj + + return dict(zip(self.__fields__, (conv(value) for value in self))) + return dict(zip(self.__fields__, self)) + + def __contains__(self, item): + if hasattr(self, "__fields__"): + return item in self.__fields__ + return super().__contains__(item) + + def __call__(self, *args): + from pyspark.sql.types import _create_row + + if len(args) > len(self): + from pyspark.errors import PySparkValueError + + raise PySparkValueError( + errorClass="TOO_MANY_VALUES", + messageParameters={ + "expected": str(len(self)), + "item": "fields", + "actual": str(len(args)), + }, + ) + return _create_row(self, args) + + def __getitem__(self, item): + if isinstance(item, (int, slice)): + return super().__getitem__(item) + from pyspark.errors import PySparkKeyError, PySparkValueError + + try: + index = self.__fields__.index(item) + return super().__getitem__(index) + except IndexError: + raise PySparkKeyError(errorClass="KEY_NOT_EXISTS", messageParameters={"key": str(item)}) + except ValueError: + raise PySparkValueError(item) + + def __getattr__(self, item): + from pyspark.errors import PySparkAttributeError + + if item.startswith("__"): + raise PySparkAttributeError( + errorClass="ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": item} + ) + try: + index = self.__fields__.index(item) + return self[index] + except IndexError: + raise PySparkAttributeError( + errorClass="ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": item} + ) + except ValueError: + raise PySparkAttributeError( + errorClass="ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": item} + ) + + def __setattr__(self, key, value): + if key != "__fields__": + from pyspark.errors import PySparkRuntimeError + + raise PySparkRuntimeError(errorClass="READ_ONLY", messageParameters={"object": "Row"}) + self.__dict__[key] = value + + def __reduce__(self): + if hasattr(self, "__fields__"): + from pyspark.sql.types import _create_row + + return _create_row, (self.__fields__, tuple(self)) + return tuple.__reduce__(self) + + def __repr__(self): + if hasattr(self, "__fields__"): + return "Row(%s)" % ", ".join( + "%s=%r" % (key, value) for key, value in zip(self.__fields__, tuple(self)) + ) + return "" % ", ".join(repr(value) for value in self) diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs index 75bd78af..29d30431 100644 --- a/crates/pyspark-rs/src/row.rs +++ b/crates/pyspark-rs/src/row.rs @@ -1,191 +1,121 @@ -//! PyO3 wrapper for spark_connect::row::Row, mirroring `pyspark.sql.Row`. +//! Conversion of native Spark values into tuple-backed Python Rows. + +use std::collections::HashMap; +use std::sync::Arc; -use pyo3::basic::CompareOp; -use pyo3::exceptions::PyIndexError; use pyo3::prelude::*; use pyo3::sync::PyOnceLock; -use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList, PySlice, PyString, PyTuple}; -use spark_connect::row::{Row, Value}; - -use crate::column::raise_pyspark; - -/// Python wrapper for a Spark Row. Behaves as `pyspark.sql.Row`, a `tuple` subclass, does, -/// though it cannot subclass `tuple` itself. -#[pyclass(name = "Row", module = "pyspark.sql.types")] -pub struct PyRow { - data: RowData, - /// Whether the row has field names (`__fields__`). As in PySpark, a row made of values - /// alone, such as `Row(1, 2)` or the row class `Row("name", "age")`, has none; the rest - /// of the client knows its fields as `_1`, `_2`, and so on. - named: bool, +use pyo3::types::{PyByteArray, PyBytes, PyDict, PyList, PyTuple}; +use spark_connect::row::{StructValue, Value}; + +/// Functions that deserialize user-defined or spatial values in each result field. +pub(crate) type FieldConverters = Arc>>>; + +static CREATE_ROW: PyOnceLock> = PyOnceLock::new(); + +/// Register the same tuple subclass under the extension and public Python imports. +/// Its source imports PySpark modules only inside methods, avoiding import cycles +/// while `pyspark.sql.types` is importing this extension. +/// The synthetic filename avoids embedding a build-machine path in distributed wheels; +/// the source is compiled into the extension and need not be installed separately. +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + let source = std::ffi::CString::new(include_str!("row.py")) + .expect("the embedded Row source contains no NUL"); + let row_module = PyModule::from_code( + module.py(), + source.as_c_str(), + c"pyspark/_row.py", + c"pyspark._row", + )?; + module.add("Row", row_module.getattr("Row")?) } -/// The Python functions that turn the values of the fields of collected rows whose types have -/// user-defined types into what PySpark gives for them, by deserializing those values; None -/// for the other fields. -pub(crate) type FieldConverters = std::sync::Arc>>>; +/// A schema path: all elements of an array or map use the same nested converter. +#[derive(Clone, Copy, Eq, Hash, PartialEq)] +enum NestedFieldPath { + Field(usize), + ArrayElement, + MapKey, + MapValue, +} -enum RowData { - /// A row that the server returned, whose Spark values become Python values when first - /// read, and stay those objects, as a PySpark row is a tuple of them: so the row equals - /// itself, with a NaN or a VARIANT value, and keeps its hash. - Spark { - row: Row, - values: PyOnceLock>, - /// The converters of its fields, when one has a user-defined type. - converters: Option, - /// Whether its binary values are `bytes`, or `bytearray`, as PySpark gives them per - /// `spark.sql.execution.pyspark.binaryAsBytes`. - binary_as_bytes: bool, - }, - /// A row made in Python, which holds its Python values as they are, as PySpark's does. - Python { - fields: Vec, - values: Py, - }, +/// Mutable names for each nested schema path, shared across the result's Arrow batches. +/// Distinct sibling structs retain independent lists even when their names are identical. +#[derive(Default)] +pub(crate) struct NestedFieldNames { + path: Vec, + fields: HashMap, Py>, } -impl PyRow { - pub fn new(row: Row) -> Self { - Self::with_converters(row, None, true) +impl NestedFieldNames { + /// Release names when an iterator closes, including when materialization fails. + pub(crate) fn clear(&mut self) { + self.path.clear(); + self.fields.clear(); } - /// A row that the server returned, whose values of user-defined types `converters` - /// deserializes when they are first read, and whose binary values are `bytes` if - /// `binary_as_bytes`, or else `bytearray`. - pub(crate) fn with_converters( - row: Row, - converters: Option, + fn value_at<'py>( + &mut self, + py: Python<'py>, + value: &Value, binary_as_bytes: bool, - ) -> Self { - PyRow { - data: RowData::Spark { - row, - values: PyOnceLock::new(), - converters, - binary_as_bytes, - }, - named: true, - } - } - - /// A row made in Python of `values`, named by `fields`. - fn of_python(fields: Vec, values: Bound<'_, PyTuple>, named: bool) -> Self { - PyRow { - data: RowData::Python { - fields, - values: values.unbind(), - }, - named, - } + path: NestedFieldPath, + ) -> PyResult> { + self.path.push(path); + let result = value_to_py(py, value, binary_as_bytes, self); + self.path.pop(); + result } - fn fields(&self) -> &[String] { - match &self.data { - RowData::Spark { row, .. } => row.fields(), - RowData::Python { fields, .. } => fields, + fn fields<'py>(&mut self, py: Python<'py>, row: &StructValue) -> PyResult> { + if let Some(fields) = self.fields.get(&self.path) { + return Ok(fields.bind(py).clone()); } + let fields = PyList::new(py, row.fields())?; + self.fields + .insert(self.path.clone(), fields.clone().unbind()); + Ok(fields) } +} - /// The values, as the tuple that a PySpark row is. - fn values<'py>(&self, py: Python<'py>) -> PyResult> { - match &self.data { - RowData::Spark { - row, - values, - converters, +/// Build a Row from named values, deserializing user-defined fields with `converters`. +/// The caller shares `fields` across result rows and nested names across the execution. +pub(crate) fn row_values_to_py( + py: Python<'_>, + fields: &Bound<'_, PyList>, + values: &[Value], + converters: Option<&FieldConverters>, + binary_as_bytes: bool, + nested_fields: &mut NestedFieldNames, +) -> PyResult> { + let values = values + .iter() + .enumerate() + .map(|(index, value)| { + let value = nested_fields.value_at( + py, + value, binary_as_bytes, - } => { - let values = values.get_or_try_init(py, || -> PyResult> { - let convert = |i: usize| { - converters - .as_ref() - .and_then(|c| c.get(i)) - .and_then(Option::as_ref) - }; - let values = row - .values() - .iter() - .enumerate() - .map(|(i, v)| { - let value = value_to_py(py, v, *binary_as_bytes)?; - match convert(i) { - Some(converter) => converter.bind(py).call1((value,)), - None => Ok(value), - } - }) - .collect::>>()?; - Ok(PyTuple::new(py, values)?.unbind()) - })?; - Ok(values.bind(py).clone()) + NestedFieldPath::Field(index), + )?; + match converters + .and_then(|converters| converters.get(index)) + .and_then(Option::as_ref) + { + Some(converter) => converter.bind(py).call1((value,)), + None => Ok(value), } - RowData::Python { values, .. } => Ok(values.bind(py).clone()), - } - } - - /// The value of the field at `index`, which is in range. - fn value<'py>(&self, py: Python<'py>, index: usize) -> PyResult> { - match &self.data { - RowData::Spark { .. } => self.values(py)?.get_item(index), - RowData::Python { values, .. } => values.bind(py).get_item(index), - } - } - - /// The value of the first field named `name`, if any. - fn value_by_name<'py>( - &self, - py: Python<'py>, - name: &str, - ) -> PyResult>> { - match self.fields().iter().position(|field| field == name) { - Some(index) => self.value(py, index).map(Some), - None => Ok(None), - } - } - - /// The field names, if the row has them (`__fields__`), and the values, as the tuple - /// that a PySpark row is, for reading the row as PySpark's schema inference does. - pub(crate) fn parts<'py>( - &self, - py: Python<'py>, - ) -> PyResult<(Option<&[String]>, Bound<'py, PyTuple>)> { - Ok((self.named.then(|| self.fields()), self.values(py)?)) - } - - /// The row with Spark values, for the rest of the client, which converts the values of a - /// row made in Python as it converts those of a tuple. - pub(crate) fn to_row(&self, py: Python<'_>) -> PyResult { - match &self.data { - RowData::Spark { row, .. } => Ok(row.clone()), - RowData::Python { fields, values } => { - let values = values - .bind(py) - .iter() - .map(|v| crate::session::py_to_value(&v)) - .collect::>>()?; - Ok(Row::new(fields.clone(), values)) - } - } - } - - /// What PySpark raises for an attribute that a row does not have. - fn no_attribute(py: Python<'_>, name: &str) -> PyErr { - raise_pyspark( - py, - "PySparkAttributeError", - "ATTRIBUTE_NOT_SUPPORTED", - &[("attr_name", name)], - ) - } - - /// The values of `other` as a tuple, if it is a row or a tuple, as a tuple operation - /// takes them. - fn tuple_of<'py>(other: &Bound<'py, PyAny>) -> PyResult>> { - if let Ok(row) = other.cast::() { - return row.borrow().values(other.py()).map(Some); - } - Ok(other.cast::().ok().cloned()) - } + }) + .collect::>>()?; + let create_row = CREATE_ROW.get_or_try_init(py, || { + py.import("pyspark.sql.types")? + .getattr("_create_row") + .map(Bound::unbind) + })?; + Ok(create_row + .bind(py) + .call1((fields, PyTuple::new(py, values)?))? + .unbind()) } /// `datetime.datetime(1970, 1, 1)`, the epoch of Spark's timestamps, without a time zone. @@ -219,10 +149,11 @@ fn micros_timedelta(py: Python<'_>, micros: i64) -> PyResult> { /// Convert a Spark `Value` into a Python object (PyO3 0.28 `IntoPyObject`), whose binary values /// are `bytes` if `binary_as_bytes`, or else `bytearray`, also in arrays, maps and structs. -pub fn value_to_py<'py>( +pub(crate) fn value_to_py<'py>( py: Python<'py>, v: &Value, binary_as_bytes: bool, + nested_fields: &mut NestedFieldNames, ) -> PyResult> { // Propagate conversion failures as Python exceptions with `?` rather than // panicking across the FFI boundary (the conversions are near-infallible, @@ -304,7 +235,12 @@ pub fn value_to_py<'py>( Value::List(items) => { let list = PyList::empty(py); for it in items { - list.append(value_to_py(py, it, binary_as_bytes)?)?; + list.append(nested_fields.value_at( + py, + it, + binary_as_bytes, + NestedFieldPath::ArrayElement, + )?)?; } list.into_any() } @@ -314,21 +250,24 @@ pub fn value_to_py<'py>( let dict = PyDict::new(py); for (k, val) in m { dict.set_item( - value_to_py(py, k, binary_as_bytes)?, - value_to_py(py, val, binary_as_bytes)?, + nested_fields.value_at(py, k, binary_as_bytes, NestedFieldPath::MapKey)?, + nested_fields.value_at(py, val, binary_as_bytes, NestedFieldPath::MapValue)?, )?; } dict.into_any() } Value::Struct(row) => { - // A struct value materializes as a Row (named fields), matching PySpark, - // rather than a plain dict. Cloning keeps shared names and independent values. - Py::new( + // A struct value materializes as a Row with the mutable names of its schema path. + let fields = nested_fields.fields(py, row)?; + row_values_to_py( py, - PyRow::with_converters(row.clone().into(), None, binary_as_bytes), + &fields, + row.values(), + None, + binary_as_bytes, + nested_fields, )? .into_bound(py) - .into_any() } Value::Variant { value, metadata } => { // A VARIANT materializes as a VariantVal (toJson/toPython decode lazily), @@ -342,307 +281,3 @@ pub fn value_to_py<'py>( } }) } - -#[pymethods] -impl PyRow { - /// Construct a Row, mirroring `pyspark.sql.Row`: `Row(name=value, ...)` makes a row with - /// field names, and `Row(v1, v2, ...)` one of values alone, which can also serve as a row - /// class: `Row("name", "age")("Alice", 11)`. - #[new] - #[pyo3(signature = (*args, **kwargs))] - fn __new__( - py: Python<'_>, - args: &Bound<'_, PyTuple>, - kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult { - let kwargs = kwargs.filter(|kw| !kw.is_empty()); - if let Some(kw) = kwargs { - if !args.is_empty() { - return Err(raise_pyspark( - py, - "PySparkValueError", - "CANNOT_SET_TOGETHER", - &[("arg_list", "args and kwargs")], - )); - } - let mut fields = Vec::with_capacity(kw.len()); - let mut values = Vec::with_capacity(kw.len()); - for (k, v) in kw.iter() { - fields.push(k.extract::()?); - values.push(v); - } - Ok(PyRow::of_python(fields, PyTuple::new(py, values)?, true)) - } else { - let fields = (0..args.len()).map(|i| format!("_{}", i + 1)).collect(); - Ok(PyRow::of_python(fields, args.clone(), false)) - } - } - - /// Number of fields. - fn __len__(&self) -> usize { - self.fields().len() - } - - fn __iter__<'py>(&self, py: Python<'py>) -> PyResult> { - self.values(py)?.try_iter() - } - - fn __reversed__<'py>(&self, py: Python<'py>) -> PyResult> { - py.import("builtins")? - .getattr("reversed")? - .call1((self.values(py)?,)) - } - - /// `x in row` tests the field names of a row that has them, and the values otherwise. - fn __contains__(&self, py: Python<'_>, item: &Bound<'_, PyAny>) -> PyResult { - if self.named { - return PyList::new(py, self.fields())?.contains(item); - } - self.values(py)?.contains(item) - } - - /// An index or a slice, as of a tuple, or a field name. - fn __getitem__<'py>( - &self, - py: Python<'py>, - key: &Bound<'py, PyAny>, - ) -> PyResult> { - if key.cast::().is_ok() { - return self.values(py)?.as_any().get_item(key); - } - if let Ok(mut i) = key.extract::() { - let n = self.fields().len() as isize; - if i < 0 { - i += n; - } - if i < 0 || i >= n { - return Err(PyIndexError::new_err("tuple index out of range")); - } - return self.value(py, i as usize); - } - // As PySpark looks the key up in `__fields__`, which a row without names lacks. - if !self.named { - return Err(Self::no_attribute(py, "__fields__")); - } - if let Ok(name) = key.cast::() { - if let Some(value) = self.value_by_name(py, &name.to_cow()?)? { - return Ok(value); - } - } - let error = py.import("pyspark.errors")?.getattr("PySparkValueError")?; - Err(PyErr::from_value(error.call1((key,))?)) - } - - /// Attribute access by field name (`row.id`). - fn __getattr__<'py>(&self, py: Python<'py>, name: &str) -> PyResult> { - if !name.starts_with("__") && self.named { - if let Some(value) = self.value_by_name(py, name)? { - return Ok(value); - } - } - Err(Self::no_attribute(py, name)) - } - - /// Only `__fields__` can be set, which names the fields, as in PySpark. - fn __setattr__( - &mut self, - py: Python<'_>, - name: &str, - value: &Bound<'_, PyAny>, - ) -> PyResult<()> { - if name != "__fields__" { - return Err(raise_pyspark( - py, - "PySparkRuntimeError", - "READ_ONLY", - &[("object", "Row")], - )); - } - let mut names = value - .try_iter()? - .map(|name| name?.str().map(|name| name.to_string())) - .collect::>>()?; - // PySpark pairs names with values as `zip` does; fields beyond the names keep theirs. - let fields = self.fields(); - names.truncate(fields.len()); - names.extend(fields[names.len()..].iter().cloned()); - match &mut self.data { - RowData::Spark { row, .. } => *row = Row::new(names, row.values().to_vec()), - RowData::Python { fields, .. } => *fields = names, - } - self.named = true; - Ok(()) - } - - /// Field names, in order: a list attribute, as PySpark's `Row.__fields__` is. A row - /// without names has none. - #[getter] - fn __fields__(&self, py: Python<'_>) -> PyResult> { - if !self.named { - return Err(Self::no_attribute(py, "__fields__")); - } - Ok(self.fields().to_vec()) - } - - /// As a dict of field name -> value (mirrors `Row.asDict`), with nested rows, also in - /// lists and dicts, as dicts if `recursive`. - #[pyo3(name = "asDict")] - #[pyo3(signature = (recursive=false))] - fn as_dict<'py>(&self, py: Python<'py>, recursive: bool) -> PyResult> { - fn conv<'py>(obj: Bound<'py, PyAny>) -> PyResult> { - let py = obj.py(); - if let Ok(row) = obj.cast::() { - return Ok(row.borrow().as_dict(py, true)?.into_any()); - } - if let Ok(list) = obj.cast::() { - let items = list.iter().map(conv).collect::>>()?; - return Ok(PyList::new(py, items)?.into_any()); - } - if let Ok(dict) = obj.cast::() { - let converted = PyDict::new(py); - for (k, v) in dict.iter() { - converted.set_item(k, conv(v)?)?; - } - return Ok(converted.into_any()); - } - Ok(obj) - } - if !self.named { - return Err(raise_pyspark( - py, - "PySparkTypeError", - "CANNOT_CONVERT_TYPE", - &[("from_type", "Row"), ("to_type", "dict")], - )); - } - let dict = PyDict::new(py); - for (name, value) in self.fields().iter().zip(self.values(py)?.iter()) { - dict.set_item(name, if recursive { conv(value)? } else { value })?; - } - Ok(dict) - } - - /// A row class, such as `Row("name", "age")`, makes a row of `args` named by its values. - #[pyo3(signature = (*args))] - fn __call__(&self, py: Python<'_>, args: &Bound<'_, PyTuple>) -> PyResult { - if args.len() > self.fields().len() { - let (expected, actual) = (self.fields().len().to_string(), args.len().to_string()); - return Err(raise_pyspark( - py, - "PySparkValueError", - "TOO_MANY_VALUES", - &[ - ("expected", expected.as_str()), - ("item", "fields"), - ("actual", actual.as_str()), - ], - )); - } - let fields = self - .values(py)? - .iter() - .take(args.len()) - .map(|name| name.str().map(|name| name.to_string())) - .collect::>>()?; - Ok(PyRow::of_python(fields, args.clone(), true)) - } - - /// Pickles as PySpark's row does: with `_create_row` from the field names and values, or - /// as `Row(*values)` for a row without names. - fn __reduce__<'py>(slf: &Bound<'py, Self>) -> PyResult> { - let py = slf.py(); - let row = slf.borrow(); - let values = row.values(py)?; - if row.named { - let create_row = py.import("pyspark.sql.types")?.getattr("_create_row")?; - let args = (PyList::new(py, row.fields())?, values).into_pyobject(py)?; - return (create_row, args).into_pyobject(py); - } - (slf.get_type(), values).into_pyobject(py) - } - - /// `Row(name='Alice', age=11)`, with the `repr()` of each value, or `` for - /// a row without names, as PySpark prints a row. - fn __repr__(&self, py: Python<'_>) -> PyResult { - let values = self.values(py)?; - let mut parts = Vec::with_capacity(values.len()); - if self.named { - for (name, value) in self.fields().iter().zip(values.iter()) { - parts.push(format!("{}={}", name, value.repr()?)); - } - return Ok(format!("Row({})", parts.join(", "))); - } - for value in values.iter() { - parts.push(value.repr()?.to_string()); - } - Ok(format!("", parts.join(", "))) - } - - /// Hashes as the tuple of its values does. - fn __hash__(&self, py: Python<'_>) -> PyResult { - self.values(py)?.hash() - } - - /// Compares as the tuple of its values does, with another row or a tuple. - fn __richcmp__<'py>( - &self, - other: &Bound<'py, PyAny>, - op: CompareOp, - ) -> PyResult> { - let py = other.py(); - match Self::tuple_of(other)? { - Some(other) => self.values(py)?.rich_compare(other, op), - None => Ok(py.NotImplemented().into_bound(py)), - } - } - - /// Concatenates as a tuple does, into a tuple. - fn __add__<'py>(&self, other: &Bound<'py, PyAny>) -> PyResult> { - let py = other.py(); - match Self::tuple_of(other)? { - Some(other) => self.values(py)?.as_any().add(other), - None => Ok(py.NotImplemented().into_bound(py)), - } - } - - fn __radd__<'py>(&self, other: &Bound<'py, PyAny>) -> PyResult> { - let py = other.py(); - match Self::tuple_of(other)? { - Some(other) => other.as_any().add(self.values(py)?), - None => Ok(py.NotImplemented().into_bound(py)), - } - } - - /// Repeats as a tuple does, into a tuple. - fn __mul__<'py>(&self, times: &Bound<'py, PyAny>) -> PyResult> { - self.values(times.py())?.as_any().mul(times) - } - - fn __rmul__<'py>(&self, times: &Bound<'py, PyAny>) -> PyResult> { - self.__mul__(times) - } - - /// Number of occurrences of `value` among the row's values (tuple.count). - fn count(&self, py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { - self.values(py)? - .as_any() - .call_method1("count", (value,))? - .extract() - } - - /// Index of the first occurrence of `value` among the row's values (tuple.index). - #[pyo3(signature = (value, *args))] - fn index( - &self, - py: Python<'_>, - value: &Bound<'_, PyAny>, - args: &Bound<'_, PyTuple>, - ) -> PyResult { - let mut call = vec![value.clone()]; - call.extend(args.iter()); - self.values(py)? - .as_any() - .call_method1("index", PyTuple::new(py, call)?)? - .extract() - } -} diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 1bfe99b1..8327dfa7 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -640,8 +640,7 @@ impl PySparkSession { } else { // A VariantVal of PySpark has a `__dict__`, so that it is taken as a row, which // PySpark's `to_row` rejects, as `row_values` does. - !(first.is_instance_of::() - || first.is_instance_of::() + !(first.is_instance_of::() || first.is_instance_of::() || first.is_instance_of::() || first.hasattr("__dict__")?) @@ -2809,19 +2808,28 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { return Ok(Value::List(items)); } if let Ok(tuple) = obj.cast::() { - let mut items = Vec::with_capacity(tuple.len()); - for item in tuple.iter() { - items.push(py_to_value(&item)?); + let mut items = tuple + .iter() + .map(|item| py_to_value(&item)) + .collect::>>()?; + // Named Rows are tuple subclasses. Keep their struct identity when they are + // nested in local data instead of treating them as plain arrays. + match obj.getattr("__fields__") { + Ok(fields) => { + let mut fields: Vec = fields.extract()?; + // Row factories may receive fewer values than names, and __fields__ + // is mutable. Pair them as schema inference does, without panicking + // in the native Row constructor when their lengths differ. + fields.truncate(items.len()); + items.truncate(fields.len()); + return Ok(Value::Struct(Row::new(fields, items).into())); + } + Err(error) if error.is_instance_of::(obj.py()) => {} + Err(error) => return Err(error), } return Ok(Value::List(items)); } - // pyspark Row -> struct: its named fields (values already core Values), so a nested - // Row inside createDataFrame data (e.g. Row(s=Row(x=1))) converts recursively. - if let Ok(pyrow) = obj.extract::>() { - return Ok(Value::Struct(pyrow.to_row(obj.py())?.into())); - } - // dict -> map, its entries in order with keys and values of their own types if let Ok(dict) = obj.cast::() { let mut map = Vec::with_capacity(dict.len()); diff --git a/crates/spark-connect/src/row.rs b/crates/spark-connect/src/row.rs index b7e95597..4e47cd20 100644 --- a/crates/spark-connect/src/row.rs +++ b/crates/spark-connect/src/row.rs @@ -425,6 +425,17 @@ impl Row { self.fields.as_slice() } + /// The immutable field-name allocation shared by rows of an Arrow record batch. + /// Language bindings use its identity to preserve per-batch field-name sharing. + /// Locally constructed rows may own their names instead. + #[doc(hidden)] + pub fn shared_fields(&self) -> Option<&Arc<[String]>> { + match &self.fields { + RowFields::Shared(fields) => Some(fields), + RowFields::Owned(_) => None, + } + } + /// Get values. pub fn values(&self) -> &[Value] { &self.values diff --git a/pyproject.toml b/pyproject.toml index e6f77f83..f6d6980a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,18 +44,24 @@ Repository = "https://github.com/apache/spark-connect-rust" "Bug Tracker" = "https://issues.apache.org/jira/browse/SPARK" [tool.ruff] -# Lint only our authored tooling/scripts. The python/pyspark tree is largely vendored +# Lint our authored tooling/scripts and the embedded Row implementation. +# The python/pyspark tree is largely vendored # from Apache Spark (pandas-on-Spark, the drop-in skin) and must stay # byte-compatible with upstream, so it is intentionally excluded from our linter. line-length = 100 target-version = "py39" extend-exclude = ["python/pyspark", "target", ".claude"] -include = ["scripts/**/*.py"] +include = ["scripts/**/*.py", "crates/pyspark-rs/src/row.py"] [tool.ruff.lint] select = ["E", "F", "W", "I", "UP", "B"] ignore = ["E501"] # line length handled by the formatter, not enforced hard +[tool.ruff.lint.per-file-ignores] +# Keep the upstream Row implementation's implicit exception context and percent +# formatting so error behavior and source comparisons stay faithful to PySpark 4.3.0. +"crates/pyspark-rs/src/row.py" = ["B904", "UP031"] + [tool.maturin] # Mixed Python/Rust layout: the pure-Python drop-in skin lives under python/pyspark, # and the PyO3 extension is placed inside it as pyspark._pyspark. diff --git a/python/.coveragerc b/python/.coveragerc index 284e3c6e..7e59e53f 100644 --- a/python/.coveragerc +++ b/python/.coveragerc @@ -3,6 +3,8 @@ # the Python glue only, excluding vendored third-party code (cloudpickle), the # test-support utilities, pandas-on-spark, generated wrappers, and the compiled # extension itself. +# The embedded crates/pyspark-rs/src/row.py executes as pyspark/_row.py, outside +# this source scope. Dedicated Row tests exercise it; Ruff checks its source. [run] source = pyspark branch = True diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py index 33a92e30..a7b6e26c 100644 --- a/python/pyspark/sql/types.py +++ b/python/pyspark/sql/types.py @@ -21,10 +21,11 @@ # Spark's nanosecond timestamp types (no PySpark 4.3.0 class; not in __all__). TimestampNTZNanosType, TimestampLTZNanosType, ) - # Row lives in the Rust core; upstream code imports it from pyspark.sql.types too. + # The extension exports the same tuple-backed Row as the public Python imports. from pyspark._pyspark import Row except ImportError: # pragma: no cover - defensive fallback when the extension is absent - # Fallback: define minimal stubs for testing purposes + # Minimal DataType stubs for import-only tests. Row and runtime conversion + # helpers require the compiled extension. class DataType: """Base class for all data types.""" def __init__(self): @@ -614,12 +615,12 @@ def convert_udt(value): fields = _udt_field_converters(dataType) if fields is None: return None - row = Row(*dataType.names) + field_names = dataType.names def convert_struct(value): if value is None: return None - return row(*[v if f is None else f(v) for f, v in zip(fields, value)]) + return _create_row(field_names, [v if f is None else f(v) for f, v in zip(fields, value)]) return convert_struct return None @@ -822,8 +823,7 @@ def _from_internal(dataType, obj): # The private helpers of PySpark's pyspark.sql.types, as of Apache Spark v4.3.0: the inference, # merging, conversion and verification of the types of Python objects, and their tables, which # PySpark Connect's modules, Apache Spark's tests and libraries import. They are PySpark's, as -# the client's types share the classes and the attributes of PySpark's; but the client's Row is -# not a tuple, so where PySpark takes a tuple of a Row, they take the client's Row too. +# the client's types share the classes and attributes of PySpark's, including its tuple Row. import ctypes # noqa: E402 import datetime # noqa: E402 import decimal # noqa: E402 diff --git a/python/tests/test_arrow_handoff.py b/python/tests/test_arrow_handoff.py index 9c85da82..47ed057d 100644 --- a/python/tests/test_arrow_handoff.py +++ b/python/tests/test_arrow_handoff.py @@ -249,6 +249,98 @@ def test_collect_reads_binary_default_with_arrow_only_schema(protocol): assert calls == {"execute": 1, "config": 1} +@pytest.mark.parametrize( + "action", + ["collect", "take", "tail", "first", "head", "head_many", "iterator", "prefetch"], +) +def test_native_result_actions_return_tuple_rows_recursively(protocol, action): + from pyspark import Row + + child_type = pa.struct([pa.field("value", pa.int64())]) + batch = pa.RecordBatch.from_arrays( + [ + pa.array([1, 2], type=pa.int64()), + pa.array([{"value": 3}, None], type=child_type), + pa.array([[{"value": 4}, None], []], type=pa.list_(child_type)), + pa.array([[("child", {"value": 5})], []], type=pa.map_(pa.string(), child_type)), + ], + names=["id", "child", "children", "by_name"], + ) + with _dataframe(protocol=protocol, batches=[batch]) as (frame, _): + if action in ("collect", "first", "head"): + result = getattr(frame, action)() + rows = result if action == "collect" else [result] + elif action in ("take", "tail"): + rows = getattr(frame, action)(2) + elif action == "head_many": + rows = frame.head(2) + else: + rows = list(frame.toLocalIterator(prefetchPartitions=action == "prefetch")) + + assert rows and all(type(row) is Row and isinstance(row, tuple) for row in rows) + row = rows[0] + assert type(row.child) is type(row.children[0]) is type(row.by_name["child"]) is Row + assert tuple.__getitem__(row.child, 0) == 3 + assert row.children[1] is None + assert row.asDict(recursive=True) == { + "id": 1, + "child": {"value": 3}, + "children": [{"value": 4}, None], + "by_name": {"child": {"value": 5}}, + } + if len(rows) == 2: + assert rows[1] == (2, None, [], {}) + + +def test_native_collect_rows_share_mutable_field_names(protocol): + batch = pa.RecordBatch.from_arrays([pa.array([1, 2])], names=["value"]) + with _dataframe(protocol=protocol, batches=[batch.slice(0, 1), batch.slice(1, 1)]) as (frame, _): + rows = frame.collect() + + assert rows[0].__fields__ is rows[1].__fields__ + rows[0].__fields__[0] = "renamed" + assert rows[1].renamed == 2 + assert rows[1].asDict() == {"renamed": 2} + + +@pytest.mark.parametrize("prefetch", [False, True]) +@pytest.mark.parametrize("one_response", [False, True]) +def test_native_iterator_rows_share_field_names_only_within_each_arrow_batch( + protocol, prefetch, one_response +): + batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3, 4])], names=["value"]) + batches = [batch.slice(0, 2), batch.slice(2, 2)] + payloads = None + if one_response: + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, batch.schema) as writer: + for part in batches: + writer.write_batch(part) + payloads = [(batch.num_rows, sink.getvalue().to_pybytes())] + with _dataframe(protocol=protocol, batches=batches, payloads=payloads) as (frame, _): + rows = list(frame.toLocalIterator(prefetchPartitions=prefetch)) + + assert rows[0].__fields__ is rows[1].__fields__ + assert rows[2].__fields__ is rows[3].__fields__ + assert rows[0].__fields__ is not rows[2].__fields__ + rows[0].__fields__[0] = "renamed" + assert rows[1].renamed == 2 + assert rows[2].value == 3 + assert rows[3].asDict() == {"value": 4} + + +def test_native_iterator_ends_after_row_materialization_fails(protocol): + # The first timestamp is year 10000, outside Python datetime's range. + batch = pa.RecordBatch.from_arrays( + [pa.array([253402300800000000, 0], type=pa.timestamp("us"))], names=["value"] + ) + with _dataframe(protocol=protocol, batches=[batch]) as (frame, _): + rows = frame.toLocalIterator() + with pytest.raises(OverflowError): + next(rows) + assert list(rows) == [] + + 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)] diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index fa5583c3..14c7dfb5 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -3887,6 +3887,243 @@ def test_row_behaves_as_a_tuple_of_its_values(): assert pickle.loads(pickle.dumps(Row(1, 2))) == (1, 2) +def test_row_is_the_public_tuple_subclass(): + from pyspark import Row as root_row + from pyspark._pyspark import Row as extension_row + from pyspark.sql import Row + from pyspark.sql.types import Row as types_row + + assert root_row is extension_row is Row is types_row + assert Row.__module__ == "pyspark.sql.types" + assert Row.__bases__ == (tuple,) + row = Row(a=1, child=Row(b=2)) + assert isinstance(row, tuple) + assert tuple.__getitem__(row, 0) == 1 + assert json.loads(json.dumps(row)) == [1, [2]] + + class DerivedRow(Row): + def total(self): + return sum(self) + + derived = DerivedRow(a=1, b=2) + assert type(derived) is DerivedRow + assert isinstance(derived, Row) and isinstance(derived, tuple) + assert derived.total() == 3 and derived.asDict() == {"a": 1, "b": 2} + + +def test_row_field_names_keep_their_identity_and_mutations(): + from pyspark.errors import PySparkKeyError + from pyspark.sql import Row + from pyspark.sql.types import _create_row + + row = Row(a=1, b=2) + fields = row.__fields__ + assert row.__fields__ is fields + fields[0] = "renamed" + assert row.renamed == 1 and "a" not in row + assert row.asDict() == {"renamed": 1, "b": 2} + assert repr(row) == "Row(renamed=1, b=2)" + fields.append("missing") + assert len(row) == 2 and tuple(row) == (1, 2) + assert "missing" in row + with pytest.raises(PySparkKeyError): + row["missing"] + + shared = ["x", "y"] + left = _create_row(shared, (1, 2)) + right = _create_row(shared, (3, 4)) + assert left.__fields__ is right.__fields__ is shared + shared[1] = "z" + assert (left.z, right.z) == (2, 4) + row.__fields__ = shared + assert row.__fields__ is shared and row.z == 2 + + factory = Row("first", "second") + made = factory(1) + assert made.__fields__ is factory + assert made.asDict() == {"first": 1} + assert repr(made) == "Row(first=1)" + + +@pytest.mark.parametrize("serializer", ["pickle", "cloudpickle"]) +def test_row_serialization_preserves_public_type_and_shared_names(serializer): + from pyspark import cloudpickle + from pyspark.sql import Row + from pyspark.sql.types import _create_row + + codec = pickle if serializer == "pickle" else cloudpickle + fields = ["value"] + rows = [_create_row(fields, (Row(child=1),)), _create_row(fields, (Row(child=2),))] + restored = codec.loads(codec.dumps(rows)) + assert all(type(row) is Row and type(row.value) is Row for row in restored) + assert restored[0].__fields__ is restored[1].__fields__ + restored[0].__fields__[0] = "renamed" + assert restored[1].renamed == Row(child=2) + unnamed = codec.loads(codec.dumps(Row(1, 2))) + assert type(unnamed) is Row and isinstance(unnamed, tuple) + assert not hasattr(unnamed, "__fields__") + assert codec.loads(codec.dumps(Row)) is Row + + +def _tuple_row_results(mock_connect, table, schema, method): + import pyarrow as pa + + spark, service = mock_connect + sink = pa.BufferOutputStream() + with pa.ipc.new_stream(sink, table.schema) as writer: + writer.write_table(table, max_chunksize=2) + service.execute_payload = message_field(7, schema) + message_field( + 2, varint_field(1, table.num_rows) + message_field(2, sink.getvalue().to_pybytes())) + frame = spark.range(table.num_rows) + return frame.collect() if method == "collect" else list( + frame.toLocalIterator(method == "prefetch")) + + +@pytest.mark.parametrize("method", ["collect", "iterator", "prefetch"]) +def test_collected_tuple_rows_share_names_by_schema_path(mock_connect, method): + import pyarrow as pa + from _mock_connect_server import data_type_proto + + # PySpark keeps top-level names per collection or iterator batch, but nested + # names belong to the schema and span batches. Equal sibling names are independent. + child = pa.struct([("v", pa.int64())]) + values = [{"v": index} for index in range(4)] + table = pa.table({ + "left": pa.array(values, child), + "right": pa.array(values, child), + "items": pa.array([[value, value] for value in values], pa.list_(child)), + "lookup": pa.array([[('a', value), ('b', value)] for value in values], + pa.map_(pa.string(), child)), + "struct_map": pa.array([[(value, value)] for value in values], pa.map_(child, child)), + }) + rows = _tuple_row_results(mock_connect, table, data_type_proto(ddl_json( + "left struct, right struct, " + "items array>, lookup map>, " + "struct_map map,struct>")), method) + assert len(rows) == 4 + for index, row in enumerate(rows): + assert type(row) is Row and isinstance(row, tuple) + assert type(row.left) is Row and row.left == (index,) + assert type(row.items[0]) is Row and row.items == [(index,), (index,)] + assert type(row.lookup['a']) is Row and row.lookup == {'a': (index,), 'b': (index,)} + assert type(row.left.__fields__) is list + assert row.left.__fields__ is rows[0].left.__fields__ + assert row.right.__fields__ is rows[0].right.__fields__ + assert row.items[0].__fields__ is row.items[1].__fields__ is rows[0].items[0].__fields__ + assert row.lookup['a'].__fields__ is row.lookup['b'].__fields__ is rows[0].lookup['a'].__fields__ + key, value = next(iter(row.struct_map.items())) + assert type(key) is type(value) is Row and key == value == (index,) + first_key, first_value = next(iter(rows[0].struct_map.items())) + assert key.__fields__ is first_key.__fields__ + assert value.__fields__ is first_value.__fields__ + assert key.__fields__ is not value.__fields__ + assert len({id(rows[0][index].__fields__) for index in (0, 1)}) == 2 + assert rows[0].left.__fields__ is not rows[0].items[0].__fields__ + assert rows[0].left.__fields__ is not rows[0].lookup['a'].__fields__ + + rows[0].left.__fields__[0] = "renamed" + rows[0].items[0].__fields__[0] = "array_value" + rows[0].lookup['a'].__fields__[0] = "map_value" + next(iter(rows[0].struct_map)).__fields__[0] = "map_key" + assert [row.left.renamed for row in rows] == list(range(4)) + assert [row.right.v for row in rows] == list(range(4)) + assert [row.items[1].array_value for row in rows] == list(range(4)) + assert [row.lookup['b'].map_value for row in rows] == list(range(4)) + assert [next(iter(row.struct_map)).map_key for row in rows] == list(range(4)) + + assert rows[0].__fields__ is rows[1].__fields__ + assert rows[2].__fields__ is rows[3].__fields__ + assert (rows[0].__fields__ is rows[2].__fields__) == (method == "collect") + rows[0].__fields__[0] = "renamed_left" + assert rows[1].renamed_left.renamed == 1 + if method != "collect": + assert rows[2].left.renamed == 2 + + +class _TupleBinaryUDT(T.UserDefinedType): + @classmethod + def sqlType(cls): + return T.StructType([T.StructField("b", T.BinaryType())]) + + @classmethod + def module(cls): + return __name__ + + def deserialize(self, datum): + return datum, datum.__fields__ + + +def _tuple_row_type_proto(data_type): + from _mock_connect_server import data_type_proto + + if isinstance(data_type, T.UserDefinedType): + value = data_type.jsonValue() + # DataType.udt=23: type=1, python_class=3, serialized_python_class=4, sql_type=5. + return message_field(23, message_field(1, "udt") + message_field(3, value["pyClass"]) + + message_field(4, value["serializedClass"]) + + message_field(5, _tuple_row_type_proto(data_type.sqlType()))) + if isinstance(data_type, T.StructType): + return message_field(21, b"".join(message_field( + 1, message_field(1, field.name) + message_field(2, _tuple_row_type_proto(field.dataType)) + + varint_field(3, int(field.nullable))) for field in data_type.fields)) + return data_type_proto(data_type.jsonValue()) + + +@pytest.mark.parametrize("method", ["collect", "iterator", "prefetch"]) +def test_tuple_rows_keep_configured_binary_and_udt_inputs(mock_connect, method): + import pyarrow as pa + + _, service = mock_connect + service.confs = {"spark.sql.execution.pyspark.binaryAsBytes": "false"} + child = pa.struct([("b", pa.binary())]) + converted = pa.struct([("b", pa.binary()), ("u", child)]) + values = [{"b": bytes([index])} for index in range(4)] + table = pa.table({ + "s": pa.array(values, child), + "converted": pa.array([{"b": value["b"], "u": value} for value in values], converted), + }) + schema = T.StructType([ + T.StructField("s", T.StructType([T.StructField("b", T.BinaryType())])), + T.StructField("converted", T.StructType([ + T.StructField("b", T.BinaryType()), T.StructField("u", _TupleBinaryUDT()), + ])), + ]) + rows = _tuple_row_results(mock_connect, table, _tuple_row_type_proto(schema), method) + assert len(rows) == 4 + for index, row in enumerate(rows): + assert type(row) is Row and type(row.s) is Row and type(row.converted) is Row + assert type(row.s.b) is type(row.converted.b) is bytearray + assert row.s.b == row.converted.b == bytes([index]) + udt_input, input_fields = row.converted.u + assert type(udt_input) is Row and isinstance(udt_input, tuple) + assert type(udt_input.b) is bytearray and udt_input.b == bytes([index]) + assert input_fields is udt_input.__fields__ and type(input_fields) is list + assert row.converted.__fields__ is rows[0].converted.__fields__ + assert type(rows[0].converted.__fields__) is list + rows[0].converted.__fields__[0] = "renamed_binary" + assert rows[-1].converted.renamed_binary == b"\x03" + + +@pytest.mark.parametrize("spatial_type", [T.GeometryType, T.GeographyType]) +def test_tuple_struct_converter_keeps_spatial_payloads_and_mutable_names(spatial_type): + from pyspark.sql.types import _udt_field_converters + + wkb = bytes.fromhex("0101000000000000000000f03f0000000000000040") + schema = T.StructType([T.StructField("s", T.StructType([ + T.StructField("b", T.BinaryType()), T.StructField("g", spatial_type(4326)), + ]))]) + convert = _udt_field_converters(schema)[0] + left = convert(Row(b=bytearray(b"x"), g=Row(wkb=wkb, srid=4326))) + right = convert(Row(b=bytearray(b"y"), g=Row(wkb=wkb, srid=4326))) + assert type(left) is Row and isinstance(left, tuple) + assert type(left.b) is bytearray and left.b == b"x" + assert type(left.g.wkb) is bytes and left.g.wkb == wkb + assert left.g == right.g + assert type(left.__fields__) is list and left.__fields__ is right.__fields__ + left.__fields__[0] = "renamed_binary" + assert right.renamed_binary == b"y" + + def test_row_fields(): from pyspark.errors import ( PySparkAttributeError, @@ -7302,7 +7539,67 @@ def test_local_input_native_numbers_inside_nested_values(mock_connect): assert dict(actual["mapping"]) == {"x": 2**63 - 1} -# Exact list/tuple conversion and pandas input IPC regressions. +@pytest.mark.parametrize("source", ["local", "pandas"]) +@pytest.mark.parametrize("make, expected", [ + (lambda: Row(a=1, b="x"), {"a": 1, "b": "x"}), + (lambda: Row("a", "b")(1, "x"), {"a": 1, "b": "x"}), + (lambda: Row("a", "b")(1), {"a": 1}), +]) +def test_tuple_rows_remain_named_structs_in_native_input(mock_connect, source, make, expected): + # A customized pandas subclass retains the native row path. Plain pandas + # frames take the separate Arrow path checked below. + value = make() + actual = _native_number_input_table( + mock_connect, [(value,)], source=source + ).column(0)[0].as_py() + assert actual == expected + + +@pytest.mark.parametrize("source", ["local", "pandas"]) +def test_unnamed_nested_row_inference_matches_pyspark(mock_connect, source): + import pyarrow as pa + + # Local inference makes positional structs; the customized pandas row path + # keeps unnamed tuple values as arrays. + column = _native_number_input_table( + mock_connect, [(Row(1, 2),)], source=source + ).column(0) + if source == "local": + assert column.type == pa.struct([pa.field("_1", pa.int64()), pa.field("_2", pa.int64())]) + assert column[0].as_py() == {"_1": 1, "_2": 2} + else: + assert pa.types.is_list(column.type) or pa.types.is_large_list(column.type) + assert column.type.value_type == pa.int64() + assert column[0].as_py() == [1, 2] + + +@pytest.mark.parametrize("value", [Row(1, 2), Row(a=1, b=2)]) +def test_plain_pandas_tuple_rows_keep_arrow_array_inference(mock_connect, value): + import pandas as pd + import pyarrow as pa + + spark, service = mock_connect + table, _ = _local_rows_relation_of( + service, spark.createDataFrame(pd.DataFrame({"value": [value]}, dtype=object))) + column = table.column(0) + assert pa.types.is_list(column.type) or pa.types.is_large_list(column.type) + assert column.type.value_type == pa.int64() + assert column[0].as_py() == [1, 2] + + +def test_pandas_tuple_field_lookup_propagates_errors(mock_connect): + class InvalidFields(tuple): + @property + def __fields__(self): + raise RuntimeError("field lookup failed") + + with pytest.raises(RuntimeError, match="^field lookup failed$"): + _native_number_input_table( + mock_connect, [(InvalidFields((1,)),)], source="pandas" + ) + + +# Exact list/tuple conversion and separate-process input IPC regressions. from datetime import date as container_input_date, datetime as container_input_datetime from decimal import Decimal as ContainerInputDecimal From 747fa1d398cd7c6baa94d3c35138d628fcb8ff79 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 22:51:38 -0800 Subject: [PATCH 12/20] [SPARK-61100][CONNECT][PYTHON] Keep successful Row field access free of error imports and preserve reference documentation --- crates/pyspark-rs/src/row.py | 100 +++++++++++++++++++++++++++++-- crates/pyspark-rs/src/row.rs | 21 ++++--- crates/pyspark-rs/src/session.rs | 4 +- crates/spark-connect/src/row.rs | 12 ++++ pyproject.toml | 6 +- 5 files changed, 121 insertions(+), 22 deletions(-) diff --git a/crates/pyspark-rs/src/row.py b/crates/pyspark-rs/src/row.py index 6abb3221..6af09bec 100644 --- a/crates/pyspark-rs/src/row.py +++ b/crates/pyspark-rs/src/row.py @@ -24,7 +24,60 @@ class Row(tuple): - """A row with optional field names, also usable as a factory of named rows.""" + """ + A row in :class:`DataFrame`. + The fields in it can be accessed: + + * like attributes (``row.key``) + * like dictionary values (``row[key]``) + + ``key in row`` will search through row keys. + + Row can be used to create a row object by using named arguments. + It is not allowed to omit a named argument to represent that the value is + None or missing. This should be explicitly set to None in this case. + + .. versionchanged:: 3.0.0 + Rows created from named arguments no longer have + field names sorted alphabetically and will be ordered in the position as + entered. + + Examples + -------- + >>> from pyspark.sql import Row + >>> row = Row(name="Alice", age=11) + >>> row + Row(name='Alice', age=11) + >>> row['name'], row['age'] + ('Alice', 11) + >>> row.name, row.age + ('Alice', 11) + >>> 'name' in row + True + >>> 'wrong_key' in row + False + + Row also can be used to create another Row like class, then it + could be used to create Row objects, such as + + >>> Person = Row("name", "age") + >>> Person + + >>> 'name' in Person + True + >>> 'wrong_key' in Person + False + >>> Person("Alice", 11) + Row(name='Alice', age=11) + + This form can also be used to create rows as tuple values, i.e. with unnamed + fields. + + >>> row1 = Row("Alice", 11) + >>> row2 = Row(name="Alice", age=11) + >>> row1 == row2 + True + """ __module__ = "pyspark.sql.types" @@ -43,7 +96,33 @@ def __new__(cls, *args, **kwargs): return tuple.__new__(cls, args) def asDict(self, recursive=False): - """Return field names and values as a dict, optionally converting nested rows.""" + """ + Return as a dict + + Parameters + ---------- + recursive : bool, optional + turns the nested Rows to dict (default: False). + + Notes + ----- + If a row contains duplicate field names, e.g., the rows of a join + between two :class:`DataFrame` that both have the fields of same names, + one of the duplicate fields will be selected by ``asDict``. ``__getitem__`` + will also return one of the duplicate fields, however returned value might + be different to ``asDict``. + + Examples + -------- + >>> from pyspark.sql import Row + >>> Row(name="Alice", age=11).asDict() == {'name': 'Alice', 'age': 11} + True + >>> row = Row(key=1, value=Row(name='a', age=2)) + >>> row.asDict() == {'key': 1, 'value': Row(name='a', age=2)} + True + >>> row.asDict(True) == {'key': 1, 'value': {'name': 'a', 'age': 2}} + True + """ if not hasattr(self, "__fields__"): from pyspark.errors import PySparkTypeError @@ -71,6 +150,7 @@ def __contains__(self, item): return super().__contains__(item) def __call__(self, *args): + """create new Row object""" from pyspark.sql.types import _create_row if len(args) > len(self): @@ -89,20 +169,22 @@ def __call__(self, *args): def __getitem__(self, item): if isinstance(item, (int, slice)): return super().__getitem__(item) - from pyspark.errors import PySparkKeyError, PySparkValueError - try: index = self.__fields__.index(item) return super().__getitem__(index) except IndexError: + from pyspark.errors import PySparkKeyError + raise PySparkKeyError(errorClass="KEY_NOT_EXISTS", messageParameters={"key": str(item)}) except ValueError: + from pyspark.errors import PySparkValueError + raise PySparkValueError(item) def __getattr__(self, item): - from pyspark.errors import PySparkAttributeError - if item.startswith("__"): + from pyspark.errors import PySparkAttributeError + raise PySparkAttributeError( errorClass="ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": item} ) @@ -110,10 +192,14 @@ def __getattr__(self, item): index = self.__fields__.index(item) return self[index] except IndexError: + from pyspark.errors import PySparkAttributeError + raise PySparkAttributeError( errorClass="ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": item} ) except ValueError: + from pyspark.errors import PySparkAttributeError + raise PySparkAttributeError( errorClass="ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": item} ) @@ -126,6 +212,7 @@ def __setattr__(self, key, value): self.__dict__[key] = value def __reduce__(self): + """Returns a tuple so Python knows how to pickle Row.""" if hasattr(self, "__fields__"): from pyspark.sql.types import _create_row @@ -133,6 +220,7 @@ def __reduce__(self): return tuple.__reduce__(self) def __repr__(self): + """Printable representation of Row used in Python REPL.""" if hasattr(self, "__fields__"): return "Row(%s)" % ", ".join( "%s=%r" % (key, value) for key, value in zip(self.__fields__, tuple(self)) diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs index 29d30431..ae5f8962 100644 --- a/crates/pyspark-rs/src/row.rs +++ b/crates/pyspark-rs/src/row.rs @@ -12,21 +12,16 @@ use spark_connect::row::{StructValue, Value}; pub(crate) type FieldConverters = Arc>>>; static CREATE_ROW: PyOnceLock> = PyOnceLock::new(); +const ROW_SOURCE: &std::ffi::CStr = pyo3::ffi::c_str!(include_str!("row.py")); -/// Register the same tuple subclass under the extension and public Python imports. +/// Expose the embedded tuple `Row` as `pyspark._pyspark.Row`, re-exported by the Python shims. /// Its source imports PySpark modules only inside methods, avoiding import cycles /// while `pyspark.sql.types` is importing this extension. /// The synthetic filename avoids embedding a build-machine path in distributed wheels; /// the source is compiled into the extension and need not be installed separately. pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - let source = std::ffi::CString::new(include_str!("row.py")) - .expect("the embedded Row source contains no NUL"); - let row_module = PyModule::from_code( - module.py(), - source.as_c_str(), - c"pyspark/_row.py", - c"pyspark._row", - )?; + let row_module = + PyModule::from_code(module.py(), ROW_SOURCE, c"pyspark/_row.py", c"pyspark._row")?; module.add("Row", row_module.getattr("Row")?) } @@ -67,7 +62,11 @@ impl NestedFieldNames { result } - fn fields<'py>(&mut self, py: Python<'py>, row: &StructValue) -> PyResult> { + fn names_at_path<'py>( + &mut self, + py: Python<'py>, + row: &StructValue, + ) -> PyResult> { if let Some(fields) = self.fields.get(&self.path) { return Ok(fields.bind(py).clone()); } @@ -258,7 +257,7 @@ pub(crate) fn value_to_py<'py>( } Value::Struct(row) => { // A struct value materializes as a Row with the mutable names of its schema path. - let fields = nested_fields.fields(py, row)?; + let fields = nested_fields.names_at_path(py, row)?; row_values_to_py( py, &fields, diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index 8327dfa7..a14f20a6 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -2812,8 +2812,8 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { .iter() .map(|item| py_to_value(&item)) .collect::>>()?; - // Named Rows are tuple subclasses. Keep their struct identity when they are - // nested in local data instead of treating them as plain arrays. + // Named Rows (tuple subclasses with `__fields__`) become structs on the generic + // object path, e.g. customized pandas object columns. Unnamed tuples stay arrays. match obj.getattr("__fields__") { Ok(fields) => { let mut fields: Vec = fields.extract()?; diff --git a/crates/spark-connect/src/row.rs b/crates/spark-connect/src/row.rs index 4e47cd20..32a0e7cd 100644 --- a/crates/spark-connect/src/row.rs +++ b/crates/spark-connect/src/row.rs @@ -501,6 +501,18 @@ mod tests { assert!(std::ptr::eq(row.fields(), cloned.fields())); } + #[test] + fn shared_fields_exposes_only_shared_name_storage() { + let fields: Arc<[String]> = vec!["id".to_string()].into(); + let row = Row::with_shared_fields(Arc::clone(&fields), vec![Value::Long(1)]); + let cloned = row.clone(); + + assert!(Arc::ptr_eq(row.shared_fields().unwrap(), &fields)); + assert!(Arc::ptr_eq(cloned.shared_fields().unwrap(), &fields)); + let owned = Row::new(vec!["id".to_string()], vec![Value::Long(1)]); + assert!(owned.shared_fields().is_none()); + } + #[test] fn shared_field_names_preserve_order_duplicates_equality_and_debug() { let names = vec![ diff --git a/pyproject.toml b/pyproject.toml index f6d6980a..0fb7b92f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,9 +45,9 @@ Repository = "https://github.com/apache/spark-connect-rust" [tool.ruff] # Lint our authored tooling/scripts and the embedded Row implementation. -# The python/pyspark tree is largely vendored -# from Apache Spark (pandas-on-Spark, the drop-in skin) and must stay -# byte-compatible with upstream, so it is intentionally excluded from our linter. +# The python/pyspark tree is largely vendored from Apache Spark (pandas-on-Spark +# and the drop-in skin). It must stay byte-compatible with upstream and is excluded +# from this linter. line-length = 100 target-version = "py39" extend-exclude = ["python/pyspark", "target", ".claude"] From b1d18b76b8d3fe66c0e27e589469b3a8d9d887dd Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 23:08:40 -0800 Subject: [PATCH 13/20] [SPARK-61100][CONNECT][PYTHON] Preserve exact tuple conversion and test Arrow-only pandas inputs --- crates/pyspark-rs/src/row.rs | 5 ++++- crates/pyspark-rs/src/session.rs | 25 +++++++++++++---------- python/tests/test_dropin_offline.py | 31 ++++++++++++++++++++++++++--- 3 files changed, 46 insertions(+), 15 deletions(-) diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs index ae5f8962..327b5707 100644 --- a/crates/pyspark-rs/src/row.rs +++ b/crates/pyspark-rs/src/row.rs @@ -25,7 +25,8 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("Row", row_module.getattr("Row")?) } -/// A schema path: all elements of an array or map use the same nested converter. +/// One step of a nested schema path. Array elements and map keys/values carry no +/// index, so all of them share the names at their path. #[derive(Clone, Copy, Eq, Hash, PartialEq)] enum NestedFieldPath { Field(usize), @@ -49,6 +50,7 @@ impl NestedFieldNames { self.fields.clear(); } + /// Convert one schema step below this path, restoring it even if conversion fails. fn value_at<'py>( &mut self, py: Python<'py>, @@ -62,6 +64,7 @@ impl NestedFieldNames { result } + /// Get mutable names for this path, using the first struct seen to initialize them. fn names_at_path<'py>( &mut self, py: Python<'py>, diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index a14f20a6..aa94d6b3 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -2814,18 +2814,21 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { .collect::>>()?; // Named Rows (tuple subclasses with `__fields__`) become structs on the generic // object path, e.g. customized pandas object columns. Unnamed tuples stay arrays. - match obj.getattr("__fields__") { - Ok(fields) => { - let mut fields: Vec = fields.extract()?; - // Row factories may receive fewer values than names, and __fields__ - // is mutable. Pair them as schema inference does, without panicking - // in the native Row constructor when their lengths differ. - fields.truncate(items.len()); - items.truncate(fields.len()); - return Ok(Value::Struct(Row::new(fields, items).into())); + if !obj.is_exact_instance_of::() { + match obj.getattr("__fields__") { + Ok(fields) => { + let mut fields: Vec = fields.extract()?; + // Row factories may receive fewer values than names, and __fields__ + // is mutable. Pair them as schema inference does, without panicking + // in the native Row constructor when their lengths differ. + fields.truncate(items.len()); + items.truncate(fields.len()); + return Ok(Value::Struct(Row::new(fields, items).into())); + } + Err(error) + if error.is_instance_of::(obj.py()) => {} + Err(error) => return Err(error), } - Err(error) if error.is_instance_of::(obj.py()) => {} - Err(error) => return Err(error), } return Ok(Value::List(items)); } diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 14c7dfb5..ba0cb1ad 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -7578,15 +7578,40 @@ def test_plain_pandas_tuple_rows_keep_arrow_array_inference(mock_connect, value) import pandas as pd import pyarrow as pa - spark, service = mock_connect - table, _ = _local_rows_relation_of( - service, spark.createDataFrame(pd.DataFrame({"value": [value]}, dtype=object))) + table = _native_number_input_table( + mock_connect, pd.DataFrame({"value": [value]}, dtype=object)) + table.validate(full=True) + assert (table.num_rows, table.num_columns) == (1, 1) + assert table.column_names == ["value"] column = table.column(0) assert pa.types.is_list(column.type) or pa.types.is_large_list(column.type) assert column.type.value_type == pa.int64() assert column[0].as_py() == [1, 2] +def test_customized_pandas_exact_nested_tuple_rows_remain_arrays(mock_connect): + import pyarrow as pa + + values = [ + ((1, None, -3), (), None, (2**63 - 1, -(2**63))), + (), + None, + ((4,), (5, 6)), + ] + table = _native_number_input_table( + mock_connect, [(value,) for value in values], source="pandas") + table.validate(full=True) + assert (table.num_rows, table.num_columns) == (4, 1) + assert table.column_names == ["value"] + assert table.column(0).type == pa.list_(pa.list_(pa.int64())) + assert table.column(0).to_pylist() == [ + [[1, None, -3], [], None, [2**63 - 1, -(2**63)]], + [], + None, + [[4], [5, 6]], + ] + + def test_pandas_tuple_field_lookup_propagates_errors(mock_connect): class InvalidFields(tuple): @property From 977fef60b4382ffae36529ce580d49687f7e1ca3 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 23:16:17 -0800 Subject: [PATCH 14/20] Clarify tuple input paths and intern Row field lookup --- crates/pyspark-rs/src/session.rs | 2 +- python/tests/test_dropin_offline.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/crates/pyspark-rs/src/session.rs b/crates/pyspark-rs/src/session.rs index aa94d6b3..f2950a9e 100644 --- a/crates/pyspark-rs/src/session.rs +++ b/crates/pyspark-rs/src/session.rs @@ -2815,7 +2815,7 @@ pub(crate) fn py_to_value(obj: &Bound<'_, PyAny>) -> PyResult { // Named Rows (tuple subclasses with `__fields__`) become structs on the generic // object path, e.g. customized pandas object columns. Unnamed tuples stay arrays. if !obj.is_exact_instance_of::() { - match obj.getattr("__fields__") { + match obj.getattr(pyo3::intern!(obj.py(), "__fields__")) { Ok(fields) => { let mut fields: Vec = fields.extract()?; // Row factories may receive fewer values than names, and __fields__ diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index ba0cb1ad..aa74eff4 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -7402,9 +7402,9 @@ def __init__(self): _local_create_dataframe([(value,)]) -# py_to_value converts the cells of object-dtype pandas input, as well as the values of -# fillna, pivot and Row conversion; local rows go through infer.convert instead, so the -# "local" source is a control for values that the exact-number fast path must not change. +# py_to_value converts object-dtype cells of customized pandas frames, plus fillna +# and pivot values; local rows use infer.convert. The "local" source below is a +# control for values that the exact-number fast path must not change. def _RowPathFrame(pd): """A subclass of pandas.DataFrame: createDataFrame converts a customized frame with its row path (py_to_value), where it converts a plain one with PySpark's Arrow conversion.""" @@ -7578,6 +7578,7 @@ def test_plain_pandas_tuple_rows_keep_arrow_array_inference(mock_connect, value) import pandas as pd import pyarrow as pa + # The default branch preserves this plain frame and its Arrow conversion route. table = _native_number_input_table( mock_connect, pd.DataFrame({"value": [value]}, dtype=object)) table.validate(full=True) From ae901b72dd1799fa0b958b75f06e11ea2cc6ec06 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 23:18:08 -0800 Subject: [PATCH 15/20] Restore pandas input regression section heading --- python/tests/test_dropin_offline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index aa74eff4..ef42cc6e 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -7625,7 +7625,7 @@ def __fields__(self): ) -# Exact list/tuple conversion and separate-process input IPC regressions. +# Exact list/tuple conversion and pandas input IPC regressions. from datetime import date as container_input_date, datetime as container_input_datetime from decimal import Decimal as ContainerInputDecimal From 13d6ee3365f7969245b006266fea1c7df664378e Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 23:20:21 -0800 Subject: [PATCH 16/20] Keep nested Row conversion helper private and document its path --- crates/pyspark-rs/src/row.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/pyspark-rs/src/row.rs b/crates/pyspark-rs/src/row.rs index 327b5707..0fbf771a 100644 --- a/crates/pyspark-rs/src/row.rs +++ b/crates/pyspark-rs/src/row.rs @@ -151,7 +151,9 @@ fn micros_timedelta(py: Python<'_>, micros: i64) -> PyResult> { /// Convert a Spark `Value` into a Python object (PyO3 0.28 `IntoPyObject`), whose binary values /// are `bytes` if `binary_as_bytes`, or else `bytearray`, also in arrays, maps and structs. -pub(crate) fn value_to_py<'py>( +/// `nested_fields` supplies shared struct names. Call through `NestedFieldNames::value_at` +/// so its schema path is set for this value and restored after conversion. +fn value_to_py<'py>( py: Python<'py>, v: &Value, binary_as_bytes: bool, From 923556609018fa87e29ddf4e9b006477c096d2aa Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 23:34:47 -0800 Subject: [PATCH 17/20] Check pandas IPC names separately from the DataFrame rename --- python/tests/test_dropin_offline.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index ef42cc6e..aca6753d 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -7583,7 +7583,13 @@ def test_plain_pandas_tuple_rows_keep_arrow_array_inference(mock_connect, value) mock_connect, pd.DataFrame({"value": [value]}, dtype=object)) table.validate(full=True) assert (table.num_rows, table.num_columns) == (1, 1) - assert table.column_names == ["value"] + # IPC keeps temporary names; the enclosing ToDF applies the public label. + assert table.column_names == ["_0"] + _, service = mock_connect + request = parse_message(service.execute_requests[-1]) + relation = parse_message(parse_message(request[3][0])[1][0]) + rename = parse_message(relation[18][0]) + assert [name.decode() for name in rename[2]] == ["value"] column = table.column(0) assert pa.types.is_list(column.type) or pa.types.is_large_list(column.type) assert column.type.value_type == pa.int64() From c3c2b9b4a1e1a1e9134ea7d1c017d66140fc9a4f Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 7 Nov 2026 23:53:15 -0800 Subject: [PATCH 18/20] Remove obsolete native Row exceptions from tuple helpers --- python/pyspark/sql/types.py | 14 ++++++-------- python/tests/test_dropin_offline.py | 8 +++----- 2 files changed, 9 insertions(+), 13 deletions(-) diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py index a7b6e26c..09f1df76 100644 --- a/python/pyspark/sql/types.py +++ b/python/pyspark/sql/types.py @@ -711,7 +711,6 @@ def _struct_to_internal(dataType, obj): if isinstance(obj, VariantVal): raise PySparkValueError("Rows cannot be of type VariantVal") names, fields = dataType.names, dataType.fields - # A Row is a tuple in PySpark. if dataType._needSerializeAnyField: # Only calling toInternal function for fields that need conversion need = dataType._needConversion @@ -720,7 +719,7 @@ def _struct_to_internal(dataType, obj): f.toInternal(obj.get(n)) if c else obj.get(n) for n, f, c in zip(names, fields, need) ) - elif isinstance(obj, (tuple, list, Row)): + elif isinstance(obj, (tuple, list)): return tuple(f.toInternal(v) if c else v for f, v, c in zip(fields, obj, need)) elif hasattr(obj, "__dict__"): d = obj.__dict__ @@ -730,7 +729,7 @@ def _struct_to_internal(dataType, obj): else: if isinstance(obj, dict): return tuple(obj.get(n) for n in names) - elif isinstance(obj, (list, tuple, Row)): + elif isinstance(obj, (list, tuple)): return tuple(obj) elif hasattr(obj, "__dict__"): d = obj.__dict__ @@ -1158,7 +1157,7 @@ def _infer_schema( if isinstance(row, dict): items = sorted(row.items()) - elif isinstance(row, (tuple, list, Row)): + elif isinstance(row, (tuple, list)): if hasattr(row, "__fields__"): # Row items = zip(row.__fields__, tuple(row)) elif hasattr(row, "_fields"): # namedtuple @@ -1369,7 +1368,7 @@ def convert_struct(obj: Any) -> Optional[Tuple]: if obj is None: return None - if isinstance(obj, (tuple, list, Row)): + if isinstance(obj, (tuple, list)): if convert_fields: return tuple(conv(v) for v, conv in zip(obj, converters)) else: @@ -1413,7 +1412,7 @@ def convert_struct(obj: Any) -> Optional[Tuple]: DayTimeIntervalType: (datetime.timedelta,), ArrayType: (list, tuple, array), MapType: (dict,), - StructType: (tuple, list, dict, Row), + StructType: (tuple, list, dict), GeometryType: (Geometry,), GeographyType: (Geography,), VariantType: ( @@ -1428,7 +1427,6 @@ def convert_struct(obj: Any) -> Optional[Tuple]: datetime.datetime, datetime.timedelta, tuple, - Row, list, dict, array, @@ -1714,7 +1712,7 @@ def verify_struct(obj: Any) -> None: if isinstance(obj, dict): for f, verifier in verifiers: verifier(obj.get(f)) - elif isinstance(obj, (tuple, list, Row)): + elif isinstance(obj, (tuple, list)): if len(obj) != len(verifiers): if name is not None: raise PySparkValueError( diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index aca6753d..045dfff8 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -248,8 +248,7 @@ def test_spatial_and_collated_json_forms_parse(): def test_private_type_helpers_are_those_of_pyspark(): # As PySpark's pyspark.sql.types: the inference, merging, conversion and verification of the - # types of Python objects, and their tables. They take the client's Row as PySpark's take - # theirs, a tuple. + # types of Python objects and their tables, including tuple-based Rows. from array import array from pyspark.errors import PySparkTypeError, PySparkValueError @@ -358,8 +357,7 @@ def __init__(self, **kwargs): def test_type_verifiers_take_and_reject_the_values_of_pyspark(): # As Apache Spark's DataTypeVerificationTests: _make_type_verifier takes the values of each - # type, of the client's Row too, and raises the errors of PySpark for the others, with the - # names of the fields. + # type, including Rows, and raises the errors of PySpark for the others with field names. from pyspark.errors import PySparkTypeError, PySparkValueError from pyspark.sql.types import _make_type_verifier @@ -383,7 +381,7 @@ def test_type_verifiers_take_and_reject_the_values_of_pyspark(): def test_types_of_python_objects_are_inferred_merged_and_converted_as_pyspark(): # As PySpark's _infer_type, _infer_schema, _merge_type and _create_converter, of the - # values of each type, also nested, of the client's Row too. + # values of each type, including nested Rows. import datetime import decimal From 8994c6841cd5baf4a3701e50f21d26b4de2dda72 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sun, 8 Nov 2026 00:19:22 -0800 Subject: [PATCH 19/20] [SPARK-61100][CONNECT][PYTHON] Omit the embedded Row synthetic path from coverage reports --- AGENTS.md | 4 ++-- python/.coveragerc | 6 ++++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 612004c3..bbe4115c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -244,8 +244,8 @@ Vendor the ones we don't have a Rust implementation for. and server-gated `tests/e2e_*.rs` (they check `std::env::var("SPARK_REMOTE")`). - **Python drop-in**: `python/.coveragerc` scopes to the hand-written shim (omits vendored/generated/Rust-only). The embedded `crates/pyspark-rs/src/row.py` - executes under a synthetic filename outside this coverage scope; its behavior is - exercised by the Row tests and its source is included in Ruff. Run: + runs as `pyspark/_row.py` without a file at that path and is explicitly omitted; + Row tests exercise its behavior and Ruff checks its source. Run: ``` cd python && PYTHONPATH=$(pwd) python -m coverage run --rcfile=.coveragerc \ -m pytest tests/test_dropin_offline.py && python -m coverage report --rcfile=.coveragerc diff --git a/python/.coveragerc b/python/.coveragerc index 7e59e53f..3004f3f1 100644 --- a/python/.coveragerc +++ b/python/.coveragerc @@ -3,8 +3,9 @@ # the Python glue only, excluding vendored third-party code (cloudpickle), the # test-support utilities, pandas-on-spark, generated wrappers, and the compiled # extension itself. -# The embedded crates/pyspark-rs/src/row.py executes as pyspark/_row.py, outside -# this source scope. Dedicated Row tests exercise it; Ruff checks its source. +# The embedded crates/pyspark-rs/src/row.py runs as pyspark/_row.py. +# That path has no file and is explicitly omitted. Row tests exercise its behavior; +# Ruff checks the embedded source. [run] source = pyspark branch = True @@ -65,6 +66,7 @@ omit = */pyspark/sql/connect/_typing.py */pyspark/sql/connect/conversion.py */pyspark/_pyspark* + */pyspark/_row.py */tests/* [report] From 6872dd36292233bdd068101419592f4a829e328e Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sun, 8 Nov 2026 00:30:45 -0800 Subject: [PATCH 20/20] Explain embedded Row coverage exclusions --- python/.coveragerc | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/.coveragerc b/python/.coveragerc index 3004f3f1..426e1c51 100644 --- a/python/.coveragerc +++ b/python/.coveragerc @@ -65,8 +65,10 @@ omit = */pyspark/sql/connect/logging.py */pyspark/sql/connect/_typing.py */pyspark/sql/connect/conversion.py + # The native extension and its embedded Row source, which has a synthetic filename. */pyspark/_pyspark* */pyspark/_row.py + # Test helpers. */tests/* [report]