diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 8df1b838..f491d237 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1404,13 +1404,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) -> 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. diff --git a/crates/pyspark-rs/src/types.rs b/crates/pyspark-rs/src/types.rs index 4eab163a..24e147c8 100644 --- a/crates/pyspark-rs/src/types.rs +++ b/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 { + // 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::()?.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). @@ -1765,12 +1776,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 { + #[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) -> PyResult { + 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::(e.to_string())) } diff --git a/crates/spark-connect/src/dataframe.rs b/crates/spark-connect/src/dataframe.rs index 8aa9db02..5edae1aa 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/crates/spark-connect/src/dataframe.rs @@ -1491,7 +1491,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(()) } diff --git a/crates/spark-connect/src/types.rs b/crates/spark-connect/src/types.rs index fd2894bb..c357273e 100644 --- a/crates/spark-connect/src/types.rs +++ b/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}; @@ -943,14 +943,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 { 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( @@ -960,33 +965,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"), 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); + } + _ => {} } } @@ -1063,6 +1120,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()) diff --git a/crates/spark-connect/tests/types_behavior.rs b/crates/spark-connect/tests/types_behavior.rs index 8f32b792..56ee941f 100644 --- a/crates/spark-connect/tests/types_behavior.rs +++ b/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; diff --git a/python/tests/_mock_connect_server.py b/python/tests/_mock_connect_server.py index ae573fad..b05469ed 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -96,6 +96,7 @@ def serve(connection): # they wait for a response. class Service: calls = 0 + execute_calls = 0 data_type = struct_schema() failure = None block_response = False @@ -104,6 +105,24 @@ class Service: # Each AnalyzePlanRequest, recorded. analyze_requests = [] + def reset(self): + # Back to the state of a new server: the tests of a module share one. + self.calls = 0 + self.execute_calls = 0 + self.data_type = struct_schema() + self.failure = None + self.block_response = False + self.response_gate.clear() + self.analyze_requests = [] + self.execute_requests = [] + self.execute_payload = b"" + self.execute_failure = None + self.execute_error = None + self.error_details = b"" + self.fetched_error_ids = [] + self.confs = {} + self.read_confs = [] + def analyze(self, request, context): self.calls += 1 self.analyze_requests.append(request) @@ -143,6 +162,7 @@ def analyze(self, request, context): read_confs = [] def execute(self, request, context): + self.execute_calls += 1 self.execute_requests.append(request) if self.execute_failure: context.abort(grpc.StatusCode.INVALID_ARGUMENT, self.execute_failure) @@ -226,6 +246,9 @@ def config(self, request, context): elif command[0] == "release": service.response_gate.set() connection.send(None) + elif command[0] == "reset": + service.reset() + connection.send(None) finally: server.stop(0).wait() @@ -249,6 +272,11 @@ def release_response(self): assert self.connection.poll(10), "Mock server did not respond" self.connection.recv() + def reset(self): + self.connection.send(("reset",)) + assert self.connection.poll(10), "Mock server did not respond" + self.connection.recv() + @contextmanager def mock_connect_session(): diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 0ff1437d..a4214c62 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -780,7 +780,7 @@ def test_structtype_object_model(): assert st.typeName() == "struct" assert st.fieldNames() == ["a", "b"] assert st.toDDL() == "a int,b string NOT NULL" - assert "root" in st.treeString() and "a: int (nullable = true)" in st.treeString() + assert "root" in st.treeString() and "a: integer (nullable = true)" in st.treeString() # toNullable makes every field nullable. assert st.toNullable().simpleString() == "struct" # json round-trips through fromJson. @@ -3465,13 +3465,20 @@ def test_create_dataframe_cannot_type_an_all_null_field_as_pyspark(data): # --------------------------------------------------------------------------- schema cache -@pytest.fixture -def mock_connect(): +@pytest.fixture(scope="module") +def _shared_mock_connect(): pytest.importorskip("grpc") with mock_connect_session() as session: yield session +@pytest.fixture +def mock_connect(_shared_mock_connect): + spark, service = _shared_mock_connect + service.reset() + return spark, service + + @pytest.mark.parametrize("first", ["schema", "columns", "dtypes"]) def test_schema_cache_one_analysis_shared_by_metadata_accessors(mock_connect, first): spark, service = mock_connect @@ -4823,3 +4830,175 @@ def __index__(self): assert _local_create_dataframe([Row(value=IndexString("x")), Row(value=5)]).isLocal() with pytest.raises(PySparkRuntimeError): _local_create_dataframe([Row(value=IndexString("x")), Row(value="y")]) + + +# --------------------------------------------------------------------------- printSchema + + +def _print_schema_field(name, data_type, nullable=True): + return ( + message_field(1, name) + + message_field(2, data_type) + + (b"\x18\x01" if nullable else b"\x18\x00") + ) + + +def _print_schema_type(): + def struct(*fields): + return message_field(21, b"".join(message_field(1, field) for field in fields)) + + integer = message_field(6, b"") + long = message_field(7, b"") + string = message_field(11, b"") + array = message_field( + 20, message_field(1, struct(_print_schema_field("item", long))) + b"\x10\x00" + ) + mapping = message_field( + 22, + message_field(1, string) + + message_field(2, struct(_print_schema_field("value", integer))) + + b"\x18\x01", + ) + return struct( + _print_schema_field("outer", struct(_print_schema_field("inner", integer, False))), + _print_schema_field("items", array, False), + _print_schema_field("lookup", mapping), + ) + + +_PRINT_SCHEMA_LEVEL1 = ( + "root\n" + " |-- outer: struct (nullable = true)\n" + " |-- items: array (nullable = false)\n" + " |-- lookup: map (nullable = true)\n\n" +) +_PRINT_SCHEMA_LEVEL2 = ( + "root\n" + " |-- outer: struct (nullable = true)\n" + " | |-- inner: integer (nullable = false)\n" + " |-- items: array (nullable = false)\n" + " | |-- element: struct (containsNull = false)\n" + " |-- lookup: map (nullable = true)\n" + " | |-- key: string\n" + " | |-- value: struct (valueContainsNull = true)\n\n" +) +_PRINT_SCHEMA_FULL = ( + "root\n" + " |-- outer: struct (nullable = true)\n" + " | |-- inner: integer (nullable = false)\n" + " |-- items: array (nullable = false)\n" + " | |-- element: struct (containsNull = false)\n" + " | | |-- item: long (nullable = true)\n" + " |-- lookup: map (nullable = true)\n" + " | |-- key: string\n" + " | |-- value: struct (valueContainsNull = true)\n" + " | | |-- value: integer (nullable = true)\n\n" +) + + +@pytest.mark.parametrize( + "level, expected", + [ + (0, _PRINT_SCHEMA_FULL), + (1, _PRINT_SCHEMA_LEVEL1), + (1.5, _PRINT_SCHEMA_LEVEL2), + (2**100, _PRINT_SCHEMA_FULL), + (float("nan"), _PRINT_SCHEMA_FULL), + ], +) +def test_print_schema_exact_python_stdout(mock_connect, capsys, level, expected): + spark, service = mock_connect + service.data_type = _print_schema_type() + df = spark.range(10) + assert df.printSchema(level) is None + captured = capsys.readouterr() + assert captured.out == expected + assert captured.err == "" + assert df.schema.treeString(maxDepth=level) == expected[:-1] + assert service.execute_calls == 0 + + +def test_print_schema_default_reuses_cached_schema(mock_connect, capsys): + spark, service = mock_connect + service.data_type = _print_schema_type() + df = spark.range(10) + for _ in range(2): + assert df.printSchema() is None + assert capsys.readouterr().out == _PRINT_SCHEMA_FULL + assert service.calls == 1 + assert service.execute_calls == 0 + + +@pytest.mark.parametrize("level", [None, []]) +def test_print_schema_falsey_level_uses_default(mock_connect, capsys, level): + spark, service = mock_connect + service.data_type = _print_schema_type() + df = spark.range(10) + assert df.printSchema(level) is None + assert capsys.readouterr().out == _PRINT_SCHEMA_FULL + with pytest.raises(TypeError): + df.schema.treeString(level) + assert service.execute_calls == 0 + + +def test_schema_tree_invalid_depth(mock_connect, capsys): + spark, service = mock_connect + service.data_type = _print_schema_type() + df = spark.range(10) + with pytest.raises(TypeError): + df.printSchema("two") + with pytest.raises(TypeError): + df.schema.treeString("two") + assert capsys.readouterr().out == "" + assert service.execute_calls == 0 + + +def test_print_schema_propagates_schema_error(mock_connect, capsys): + spark, service = mock_connect + service.failure = "rpc" + with pytest.raises(Exception, match="schema unavailable"): + spark.range(10).printSchema() + assert capsys.readouterr().out == "" + assert service.execute_calls == 0 + + +def test_schema_tree_type_names_and_escaped_field_names(): + schema = T.StructType( + [ + T.StructField("é🌱\\n\n\r\t\f\b\v\a", T.LongType(), False), + T.StructField("decimal", T.DecimalType(10, 2)), + T.StructField("small", T.ByteType()), + T.StructField("short", T.ShortType()), + ] + ) + expected = ( + "root\n" + " |-- é🌱\\n\\n\\r\\t\\f\\b\\v\\a: long (nullable = false)\n" + " |-- decimal: decimal(10,2) (nullable = true)\n" + " |-- small: byte (nullable = true)\n" + " |-- short: short (nullable = true)\n" + ) + assert schema.treeString() == expected + assert schema.treeString(0) == expected + assert schema.treeString(-1) == expected + + +@pytest.mark.parametrize( + "python_class, expected", + [ + ("pyspark.ml.linalg.VectorUDT", "vectorudt"), + (None, "udt"), + ], +) +def test_print_schema_udt_class_name(mock_connect, capsys, python_class, expected): + spark, service = mock_connect + udt = message_field(1, "udt") + if python_class is not None: + udt += message_field(3, python_class) + service.data_type = message_field( + 21, message_field(1, _print_schema_field("features", message_field(23, udt))) + ) + assert spark.range(10).printSchema() is None + assert capsys.readouterr().out == f"root\n |-- features: {expected} (nullable = true)\n\n" + assert service.calls == 1 + assert service.execute_calls == 0