[SPARK-60695][CONNECT] Print DataFrame schemas without executing queries

crates/pyspark-rs/src/dataframe.rs

@@ -1359,13 +1359,23 @@ impl PyDataFrame {
         self.dataframe.is_empty().to_pyerr()
     }
 
-    /// Print the schema of the DataFrame. `level` (max nesting depth to print) is
-    /// accepted for signature parity with reference pyspark; the connect core does not
-    /// yet thread a depth into the schema tree-string, so it is currently ignored.
+    /// Print the schema through Python stdout, honoring the maximum nesting depth.
     #[pyo3(signature = (level=None))]
-    #[allow(unused_variables)]
-    fn printSchema(&self, level: Option<i32>) -> PyResult<()> {
-        self.dataframe.show(0).to_pyerr()
+    fn printSchema(&self, py: Python<'_>, level: Option<&Bound<'_, PyAny>>) -> PyResult<()> {
+        // Reference's `if level:` selects the default for false values before
+        // resolving the schema; positive fractional depths behave as ceil(level).
+        let level = match level {
+            Some(level) if level.is_truthy()? => Some(level),
+            _ => None,
+        };
+        let schema = py.detach(|| self.dataframe.schema().to_pyerr())?;
+        let depth = level
+            .map(crate::types::normalize_tree_depth)
+            .transpose()?
+            .unwrap_or(i32::MAX);
+        let tree = schema.tree_string_with_depth(depth).to_pyerr()?;
+        py.import("builtins")?.getattr("print")?.call1((tree,))?;
+        Ok(())
     }
 
     /// Get the schema as a list of (name, type) tuples.

crates/pyspark-rs/src/types.rs

@@ -4,6 +4,17 @@ use pyo3::prelude::*;
 use spark_connect::types::{DataType, StructField};
 use std::collections::BTreeMap;
 
+/// Normalize the depth used by both Python schema-tree entry points.
+pub(crate) fn normalize_tree_depth(depth: &Bound<'_, PyAny>) -> PyResult<i32> {
+    // Reference compares > 0 at each level, so positive fractional depths
+    // behave as ceil(depth). Nonpositive depths and NaN print the full tree.
+    if !depth.gt(0)? || depth.ge(i32::MAX)? {
+        Ok(i32::MAX)
+    } else {
+        Ok(depth.extract::<f64>()?.ceil() as i32)
+    }
+}
+
 /// Make a DataType picklable so it round-trips to an official-pyspark UDF worker:
 /// pickle reconstructs it via `pyspark.sql.types._parse_datatype_json_string(json)`
 /// (session-free, present in both our shim and the worker's official pyspark).
@@ -1810,12 +1821,17 @@ impl PyStructType {
 
     /// Tree string. Mirrors `StructType.treeString(maxDepth)`.
     #[allow(non_snake_case)]
-    #[pyo3(name = "treeString", signature = (maxDepth = 2147483647))]
-    fn __st_tree_string(&self, maxDepth: i32) -> PyResult<String> {
+    #[pyo3(
+        name = "treeString",
+        signature = (maxDepth = Python::attach(|py| pyo3::types::PyInt::new(py, i32::MAX).into_any().unbind())),
+        text_signature = "($self, maxDepth=2147483647)"
+    )]
+    fn __st_tree_string(&self, py: Python<'_>, maxDepth: Py<PyAny>) -> PyResult<String> {
+        let depth = normalize_tree_depth(maxDepth.bind(py))?;
         DataType::Struct {
             fields: self.fields.clone(),
         }
-        .tree_string_with_depth(maxDepth)
+        .tree_string_with_depth(depth)
         .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
     }
 

crates/spark-connect/src/dataframe.rs

@@ -1499,7 +1499,7 @@ impl DataFrame {
     /// Print the schema of this DataFrame.
     pub fn print_schema(&self) -> Result<()> {
         let schema = self.schema()?;
-        println!("{}", schema);
+        print!("{}", schema.tree_string()?);
         Ok(())
     }
 

crates/spark-connect/src/types.rs

@@ -4,7 +4,7 @@
 //! and the Spark Connect protobuf representation.
 
 use std::collections::BTreeMap;
-use std::fmt;
+use std::fmt::{self, Write};
 use std::hash::{Hash, Hasher};
 
 use spark_connect_core::error::{Result, SparkError};
@@ -858,14 +858,19 @@ impl DataType {
         self.tree_string_with_depth(i32::MAX)
     }
 
-    /// Like [`DataType::tree_string`], but stops recursing into nested structs
-    /// once `max_depth` nesting levels have been printed (top-level fields are
-    /// depth 1). Mirrors `StructType.treeString(maxDepth)`.
+    /// Like [`DataType::tree_string`], but limits nesting through structs,
+    /// arrays, and maps to `max_depth` levels. Top-level fields are depth 1;
+    /// zero or negative values print the full tree.
+    /// Mirrors `StructType.treeString(maxDepth)`.
     pub fn tree_string_with_depth(&self, max_depth: i32) -> Result<String> {
         match self {
             DataType::Struct { .. } => {
                 let mut out = String::from("root\n");
-                self.append_tree(&mut out, " |", 1, max_depth);
+                self.append_tree(
+                    &mut out,
+                    " |",
+                    if max_depth > 0 { max_depth } else { i32::MAX },
+                );
                 Ok(out)
             }
             _ => Err(SparkError::value(
@@ -875,33 +880,85 @@ impl DataType {
         }
     }
 
-    fn append_tree(&self, out: &mut String, prefix: &str, depth: i32, max_depth: i32) {
-        if let DataType::Struct { fields } = self {
-            for f in fields {
-                // A nested struct is rendered as the bare type name ("struct") and its
-                // fields are shown by recursing below, so `max_depth` can truncate them.
-                // Using `simple_string()` here would inline the children (e.g.
-                // "struct<inner:int>"), leaking them past `max_depth`. Non-struct types
-                // have no depth-controlled children, so their simple string is used.
-                let is_struct = matches!(f.data_type, DataType::Struct { .. });
-                let type_repr = if is_struct {
-                    f.data_type.type_name()
-                } else {
-                    f.data_type.simple_string()
-                };
-                out.push_str(&format!(
-                    "{}-- {}: {} (nullable = {})\n",
-                    prefix, f.name, type_repr, f.nullable
-                ));
-                if is_struct && depth < max_depth {
-                    f.data_type.append_tree(
+    fn tree_type_name(&self) -> String {
+        match self {
+            DataType::Decimal { .. }
+            | DataType::Char { .. }
+            | DataType::Varchar { .. }
+            | DataType::Time { .. }
+            | DataType::YearMonthInterval { .. }
+            | DataType::DayTimeInterval { .. } => self.simple_string(),
+            DataType::Udt {
+                python_class: Some(python_class),
+                ..
+            } => python_class
+                .rsplit('.')
+                .next()
+                .filter(|class_name| !class_name.is_empty())
+                // UserDefinedType.typeName() preserves the entire class name,
+                // unlike DataType.typeName(), which strips the "Type" suffix.
+                .map(str::to_lowercase)
+                .unwrap_or_else(|| self.type_name()),
+            _ => self.type_name(),
+        }
+    }
+
+    fn append_tree(&self, out: &mut String, prefix: &str, remaining: i32) {
+        if remaining <= 0 {
+            return;
+        }
+        match self {
+            DataType::Struct { fields } => {
+                let child_prefix = format!("{prefix}    |");
+                for f in fields {
+                    out.push_str(prefix);
+                    out.push_str("-- ");
+                    push_escaped_name(out, &f.name);
+                    writeln!(
                         out,
-                        &format!("{}    |", prefix.trim_end_matches('|')),
-                        depth + 1,
-                        max_depth,
-                    );
+                        ": {} (nullable = {})",
+                        f.data_type.tree_type_name(),
+                        f.nullable
+                    )
+                    .expect("writing to a String cannot fail");
+                    f.data_type.append_tree(out, &child_prefix, remaining - 1);
                 }
             }
+            DataType::Array {
+                element_type,
+                contains_null,
+            } => {
+                let child_prefix = format!("{prefix}    |");
+                writeln!(
+                    out,
+                    "{}-- element: {} (containsNull = {})",
+                    prefix,
+                    element_type.tree_type_name(),
+                    contains_null
+                )
+                .expect("writing to a String cannot fail");
+                element_type.append_tree(out, &child_prefix, remaining - 1);
+            }
+            DataType::Map {
+                key_type,
+                value_type,
+                value_contains_null,
+            } => {
+                let child_prefix = format!("{prefix}    |");
+                writeln!(out, "{}-- key: {}", prefix, key_type.tree_type_name())
+                    .expect("writing to a String cannot fail");
+                key_type.append_tree(out, &child_prefix, remaining - 1);
+                writeln!(
+                    out,
+                    "{}-- value: {} (valueContainsNull = {})",
+                    prefix,
+                    value_type.tree_type_name(),
+                    value_contains_null
+                )
+                .expect("writing to a String cannot fail");
+                value_type.append_tree(out, &child_prefix, remaining - 1);
+            }
+            _ => {}
         }
     }
 
@@ -978,6 +1035,21 @@ impl DataType {
     }
 }
 
+fn push_escaped_name(out: &mut String, name: &str) {
+    for ch in name.chars() {
+        match ch {
+            '\n' => out.push_str("\\n"),
+            '\r' => out.push_str("\\r"),
+            '\t' => out.push_str("\\t"),
+            '\x0c' => out.push_str("\\f"),
+            '\x08' => out.push_str("\\b"),
+            '\x0b' => out.push_str("\\v"),
+            '\x07' => out.push_str("\\a"),
+            _ => out.push(ch),
+        }
+    }
+}
+
 impl fmt::Display for DataType {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "{}", self.simple_string())

crates/spark-connect/tests/types_behavior.rs

@@ -280,7 +280,7 @@ fn struct_to_ddl_tree_and_nullable() {
     assert_eq!(st.to_ddl().unwrap(), "a int,b string NOT NULL");
     let tree = st.tree_string().unwrap();
     assert!(tree.starts_with("root\n"));
-    assert!(tree.contains("a: int (nullable = true)"));
+    assert!(tree.contains("a: integer (nullable = true)"));
     assert!(tree.contains("b: string (nullable = false)"));
     // toNullable makes every field nullable.
     if let DataType::Struct { fields } = st.to_nullable() {
@@ -314,13 +314,93 @@ fn tree_string_with_depth_truncates_nested_structs() {
     // Full depth shows the nested field, rendering the struct as the bare type name.
     let full = nested.tree_string().unwrap();
     assert!(full.contains("outer: struct (nullable = true)"));
-    assert!(full.contains("inner: int (nullable = true)"));
+    assert!(full.contains("inner: integer (nullable = true)"));
     // max_depth = 1 prints only the top level; the nested field is omitted.
     let shallow = nested.tree_string_with_depth(1).unwrap();
     assert!(shallow.contains("outer: struct (nullable = true)"));
     assert!(!shallow.contains("inner"));
 }
 
+#[test]
+fn tree_string_uses_python_udt_class_names() {
+    use spark_connect::types::StructField;
+    let vector_udt = || DataType::Udt {
+        type_str: "udt".to_string(),
+        jvm_class: None,
+        python_class: Some("pyspark.ml.linalg.VectorUDT".to_string()),
+        serialized_python_class: None,
+        sql_type: Some(Box::new(DataType::Integer)),
+    };
+    let schema = DataType::Struct {
+        fields: vec![
+            StructField {
+                name: "features".to_string(),
+                data_type: vector_udt(),
+                nullable: true,
+                metadata: Default::default(),
+            },
+            StructField {
+                name: "array".to_string(),
+                data_type: DataType::Array {
+                    element_type: Box::new(vector_udt()),
+                    contains_null: false,
+                },
+                nullable: false,
+                metadata: Default::default(),
+            },
+            StructField {
+                name: "map".to_string(),
+                data_type: DataType::Map {
+                    key_type: Box::new(DataType::String {
+                        collation: String::new(),
+                    }),
+                    value_type: Box::new(vector_udt()),
+                    value_contains_null: true,
+                },
+                nullable: true,
+                metadata: Default::default(),
+            },
+        ],
+    };
+
+    assert_eq!(
+        schema.tree_string().unwrap(),
+        concat!(
+            "root\n",
+            " |-- features: vectorudt (nullable = true)\n",
+            " |-- array: array (nullable = false)\n",
+            " |    |-- element: vectorudt (containsNull = false)\n",
+            " |-- map: map (nullable = true)\n",
+            " |    |-- key: string\n",
+            " |    |-- value: vectorudt (valueContainsNull = true)\n",
+        )
+    );
+}
+
+#[test]
+fn tree_string_keeps_generic_udt_name_without_python_class() {
+    use spark_connect::types::StructField;
+    let schema = DataType::Struct {
+        fields: vec![StructField {
+            name: "features".to_string(),
+            data_type: DataType::Udt {
+                type_str: "udt".to_string(),
+                jvm_class: Some("org.apache.spark.ml.linalg.VectorUDT".to_string()),
+                python_class: None,
+                serialized_python_class: None,
+                sql_type: Some(Box::new(DataType::Integer)),
+            },
+            nullable: true,
+            metadata: Default::default(),
+        }],
+    };
+
+    assert_eq!(
+        schema.tree_string().unwrap(),
+        "root\n |-- features: udt (nullable = true)\n",
+    );
+}
+
 #[test]
 fn from_json_str_roundtrips() {
     use spark_connect::types::StructField;