[SPARK-61100][CONNECT][PYTHON] Preserve tuple semantics for Python Row

.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

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`
+  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
@@ -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;

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<Vec<PyRow>> {
+    fn collect(&self, py: Python<'_>) -> PyResult<Vec<Py<PyAny>>> {
         // 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<Option<PyRow>> {
+    fn first(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
         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<Vec<PyRow>> {
+    fn take(&self, py: Python<'_>, num: usize) -> PyResult<Vec<Py<PyAny>>> {
         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<Vec<PyRow>> {
+    fn tail(&self, py: Python<'_>, num: i32) -> PyResult<Vec<Py<PyAny>>> {
         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,
@@ -4089,47 +4088,50 @@ fn check_collectable(
     Ok(())
 }
 
-/// Whether the binary values of a result of `schema` of `dataframe` are `bytes`, as PySpark
-/// Connect's `DataFrame._get_binary_as_bytes` gives it: the conf
-/// `spark.sql.execution.pyspark.binaryAsBytes`, true unless it is set to another value than
-/// `true`. The conf is read only for a result with a binary value, the only one it changes.
-fn binary_as_bytes(
-    py: Python<'_>,
-    dataframe: &DataFrame,
-    schema: Option<&spark_connect::types::DataType>,
-) -> PyResult<bool> {
-    let binary =
-        |t: &spark_connect::types::DataType| matches!(t, spark_connect::types::DataType::Binary);
-    if !schema.is_some_and(|schema| has_type(schema, &binary)) {
-        return Ok(true);
-    }
-    let value = py
-        .detach(|| {
-            dataframe
-                .spark_session()
-                .conf()
-                .get_with_default("spark.sql.execution.pyspark.binaryAsBytes", Some("true"))
-        })
-        .to_pyerr()?;
+/// The binary result setting, as PySpark's `DataFrame._get_binary_as_bytes` reads it,
+/// including when the server sends only an Arrow schema or the result has no binary fields.
+/// Collect reads it after execution; a local iterator snapshots it before its first fetch.
+fn read_binary_as_bytes(py: Python<'_>, dataframe: &DataFrame) -> PyResult<bool> {
+    let value = detached(py, || {
+        dataframe
+            .spark_session()
+            .conf()
+            .get_with_default("spark.sql.execution.pyspark.binaryAsBytes", Some("true"))
+            .map_err(Box::new)
+    })?;
     Ok(value.is_some_and(|value| value.eq_ignore_ascii_case("true")))
 }
 
-/// The collected `rows` of a result of `schema` of `dataframe`, whose values of user-defined
-/// types are deserialized when first read.
+/// 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<spark_connect::row::Row>,
+    rows: &[spark_connect::row::Row],
     schema: Option<&spark_connect::types::DataType>,
-) -> PyResult<Vec<PyRow>> {
+) -> PyResult<Vec<Py<PyAny>>> {
     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))
-        .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.
@@ -4146,8 +4148,12 @@ pub struct PyLocalRowIterator {
     /// The converters of the fields of user-defined types of the rows, once the schema of the
     /// result is known: with the first row.
     converters: Option<Option<FieldConverters>>,
-    /// Whether the binary values of the rows are `bytes`, known with the converters.
-    binary_as_bytes: bool,
+    /// Snapshotted before the first fetch, as the body of PySpark's generator does.
+    binary_as_bytes: Option<bool>,
+    /// Top-level names shared only by rows of the current Arrow record batch.
+    fields: Option<(std::sync::Arc<[String]>, Py<PyList>)>,
+    /// Nested names belong to schema paths and stay shared across iterator batches.
+    nested_fields: NestedFieldNames,
 }
 
 impl PyLocalRowIterator {
@@ -4158,7 +4164,9 @@ impl PyLocalRowIterator {
             iterator: None,
             finished: false,
             converters: None,
-            binary_as_bytes: true,
+            binary_as_bytes: None,
+            fields: None,
+            nested_fields: NestedFieldNames::default(),
         }
     }
 
@@ -4172,6 +4180,61 @@ impl PyLocalRowIterator {
         }
         self.iterator.as_mut()?.next()
     }
+
+    /// Prepare and fetch one Python row; the public wrapper closes on exhaustion or error.
+    fn next_python_row(&mut self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
+        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();
+                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),
+        }
+    }
 }
 
 #[pymethods]
@@ -4186,47 +4249,24 @@ 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));
     }
 
     /// Fetch the next row from the iterator.
     ///
-    /// Releases the GIL during the blocking fetch to allow other threads to run.
-    fn __next__(&mut self, py: Python<'_>) -> PyResult<Option<PyRow>> {
+    /// The first call snapshots the binary result setting before starting execution.
+    /// Both the Config RPC and blocking row fetch release the GIL.
+    fn __next__(&mut self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
         if self.finished {
             return Ok(None);
         }
-        // Release the GIL while fetching from the Rust iterator (which may block on network I/O).
-        let next = py.detach(|| self.next_row());
-        if self.converters.is_none() {
-            if let Some(schema) = self.iterator.as_ref().and_then(LocalRowIterator::schema) {
-                // The result is known once its first message arrives, possibly with no rows.
-                if let Err(error) = check_collectable(py, Some(&schema)) {
-                    self.finished = true;
-                    return Err(error);
-                }
-                let converters = udt_field_converters(py, Some(&schema))?;
-                self.binary_as_bytes = binary_as_bytes(py, &self.dataframe, Some(&schema))?;
-                self.converters = Some(converters);
-            }
-        }
-        match next {
-            Some(Ok(row)) => {
-                let converters = self.converters.clone().flatten();
-                Ok(Some(PyRow::with_converters(
-                    row,
-                    converters,
-                    self.binary_as_bytes,
-                )))
-            }
-            // As PySpark raises the error of the execution, which ends the iteration.
-            Some(Err(error)) => {
-                self.finished = true;
-                Err(crate::errors::spark_error_to_py_exception(error))
-            }
-            None => {
-                self.finished = true;
-                Ok(None)
+        match self.next_python_row(py) {
+            Ok(Some(row)) => Ok(Some(row)),
+            result => {
+                self.close(py);
+                result
             }
         }
     }