From 616c42b1748507c8e9221203c6b4331a1ea2dcb1 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Thu, 29 Oct 2026 00:52:35 -0700 Subject: [PATCH 1/4] [SPARK-60695][CONNECT] Print DataFrame schemas without executing queries --- crates/pyspark-rs/src/dataframe.rs | 24 +- crates/spark-connect/src/types.rs | 98 ++++++-- crates/spark-connect/tests/types_behavior.rs | 4 +- python/tests/test_dropin_offline.py | 232 ++++++++++++++++++- 4 files changed, 324 insertions(+), 34 deletions(-) diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index d74891d7..540d93e2 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1359,13 +1359,25 @@ 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<()> { + let schema = py.detach(|| self.dataframe.schema()).to_pyerr()?; + // Reference PySpark treats false values and nonpositive depths as unlimited. + // Positive fractional depths include the partially remaining level as well. + let depth = match level { + Some(level) if level.is_truthy()? && level.gt(0)? => { + if level.ge(i32::MAX)? { + i32::MAX + } else { + level.extract::()?.ceil() as i32 + } + } + _ => 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/spark-connect/src/types.rs b/crates/spark-connect/src/types.rs index 1f06c73f..5165b91d 100644 --- a/crates/spark-connect/src/types.rs +++ b/crates/spark-connect/src/types.rs @@ -865,7 +865,11 @@ impl DataType { 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 +879,77 @@ 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() - }; + fn tree_type_name(&self) -> String { + match self { + DataType::Decimal { .. } + | DataType::Char { .. } + | DataType::Varchar { .. } + | DataType::Time { .. } + | DataType::YearMonthInterval { .. } + | DataType::DayTimeInterval { .. } => self.simple_string(), + _ => self.type_name(), + } + } + + fn append_tree(&self, out: &mut String, prefix: &str, remaining: i32) { + if remaining <= 0 { + return; + } + let child_prefix = format!("{prefix} |"); + match self { + DataType::Struct { fields } => { + for f in fields { + let name = f + .name + .replace('\n', "\\n") + .replace('\r', "\\r") + .replace('\t', "\\t") + .replace('\x0c', "\\f") + .replace('\x08', "\\b") + .replace('\x0b', "\\v") + .replace('\x07', "\\a"); + out.push_str(&format!( + "{}-- {}: {} (nullable = {})\n", + prefix, + name, + f.data_type.tree_type_name(), + f.nullable + )); + f.data_type.append_tree(out, &child_prefix, remaining - 1); + } + } + DataType::Array { + element_type, + contains_null, + } => { out.push_str(&format!( - "{}-- {}: {} (nullable = {})\n", - prefix, f.name, type_repr, f.nullable + "{}-- element: {} (containsNull = {})\n", + prefix, + element_type.tree_type_name(), + contains_null )); - if is_struct && depth < max_depth { - f.data_type.append_tree( - out, - &format!("{} |", prefix.trim_end_matches('|')), - depth + 1, - max_depth, - ); - } + element_type.append_tree(out, &child_prefix, remaining - 1); } + DataType::Map { + key_type, + value_type, + value_contains_null, + } => { + out.push_str(&format!( + "{}-- key: {}\n", + prefix, + key_type.tree_type_name() + )); + key_type.append_tree(out, &child_prefix, remaining - 1); + out.push_str(&format!( + "{}-- value: {} (valueContainsNull = {})\n", + prefix, + value_type.tree_type_name(), + value_contains_null + )); + value_type.append_tree(out, &child_prefix, remaining - 1); + } + _ => {} } } diff --git a/crates/spark-connect/tests/types_behavior.rs b/crates/spark-connect/tests/types_behavior.rs index 8f32b792..0b5aff09 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,7 +314,7 @@ 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)")); diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 9988f655..a7b92d0d 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -416,7 +416,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. @@ -3085,3 +3085,233 @@ def allow_response_after_request(): else: assert result == [("id", "bigint")] assert service.calls == 1 + + +# printSchema uses a separate mock process so the real transport is exercised. +def _print_schema_varint(value): + result = bytearray() + while value > 127: + result.append((value & 127) | 128) + value >>= 7 + result.append(value) + return bytes(result) + + +def _print_schema_message(field, value): + if isinstance(value, str): + value = value.encode() + return _print_schema_varint(field * 8 + 2) + _print_schema_varint(len(value)) + value + + +def _print_schema_type(): + msg = _print_schema_message + field = _print_schema_field + integer = msg(6, b"") + long = msg(7, b"") + string = msg(11, b"") + struct = lambda *fields: msg(21, b"".join(msg(1, f) for f in fields)) + array = msg(20, msg(1, struct(field("item", long))) + b"\x10\x00") + mapping = msg(22, msg(1, string) + msg(2, struct(field("value", integer))) + b"\x18\x01") + return struct(field("outer", struct(field("inner", integer, False))), + field("items", array, False), field("lookup", mapping)) + + +def _print_schema_field(name, data_type, nullable=True): + return (_print_schema_message(1, name) + _print_schema_message(2, data_type) + (b"\x18\x01" if nullable else b"\x18\x00") + + _print_schema_message(4, '{"note":"original"}')) + + +def _print_schema_serve(connection): + import grpc + from concurrent.futures import ThreadPoolExecutor + # A separate process also exercises columns/dtypes, whose first access may + # hold the client's GIL while waiting for the RPC. + class Service: + calls = 0 + execute_calls = 0 + data_type = _print_schema_type() + failure = None + + def analyze(self, request, context): + self.calls += 1 + if self.failure == "rpc": + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "schema unavailable") + if self.failure == "no-result": + return b"" + if self.failure == "missing-schema": + return _print_schema_message(2, b"") + if self.failure == "invalid-type": + return _print_schema_message(2, _print_schema_message(1, b"")) + # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1. + # session_id is the first length-delimited field in these requests. + assert request[0] == 10 and request[1] < 128 + session_id = request[2:2 + request[1]] + return (_print_schema_message(1, session_id) + + _print_schema_message(2, _print_schema_message(1, self.data_type))) + + def execute(self, request, context): + self.execute_calls += 1 + context.abort(grpc.StatusCode.INTERNAL, "printSchema must not execute a plan") + yield b"" + + service = Service() + with ThreadPoolExecutor(max_workers=4) as executor: + server = grpc.server(executor) + handler = grpc.unary_unary_rpc_method_handler(service.analyze) + server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( + "spark.connect.SparkConnectService", {"AnalyzePlan": handler, "ExecutePlan": grpc.unary_stream_rpc_method_handler(service.execute), + "ReleaseSession": grpc.unary_unary_rpc_method_handler(lambda request, context: request[:2 + request[1]])}),)) + port = server.add_insecure_port("127.0.0.1:0") + server.start() + connection.send(port) + try: + while True: + command = connection.recv() + if command[0] == "stop": + break + if command[0] == "get": + connection.send(getattr(service, command[1])) + elif command[0] == "set": + setattr(service, command[1], command[2]) + connection.send(None) + finally: + server.stop(0).wait() + + +class _PrintSchemaService: + def __init__(self, connection): + object.__setattr__(self, "connection", connection) + + def __getattr__(self, name): + self.connection.send(("get", name)) + assert self.connection.poll(10), "Mock server did not respond" + return self.connection.recv() + + def __setattr__(self, name, value): + self.connection.send(("set", name, value)) + assert self.connection.poll(10), "Mock server did not respond" + self.connection.recv() + + +@pytest.fixture +def print_schema_service(): + import multiprocessing + pytest.importorskip("grpc") + context = multiprocessing.get_context("spawn") + parent, child = context.Pipe() + process = context.Process(target=_print_schema_serve, args=(child,)) + spark = None + process.start() + try: + assert parent.poll(20), "Mock server did not start" + port = parent.recv() + from pyspark.sql import SparkSession + spark = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").create() + yield spark, _PrintSchemaService(parent) + finally: + if spark is not None: + spark.stop() + if process.is_alive(): + parent.send(("stop",)) + process.join(10) + if process.is_alive(): + process.terminate() + process.join(10) + parent.close() + child.close() + + +_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", [ + (None, _PRINT_SCHEMA_FULL), (0, _PRINT_SCHEMA_FULL), + (-1, _PRINT_SCHEMA_FULL), (-10, _PRINT_SCHEMA_FULL), + (1, _PRINT_SCHEMA_LEVEL1), (2, _PRINT_SCHEMA_LEVEL2), + (3, _PRINT_SCHEMA_FULL), (2**100, _PRINT_SCHEMA_FULL), + (False, _PRINT_SCHEMA_FULL), (True, _PRINT_SCHEMA_LEVEL1), + (0.5, _PRINT_SCHEMA_LEVEL1), (1.5, _PRINT_SCHEMA_LEVEL2), + (float("inf"), _PRINT_SCHEMA_FULL), (float("nan"), _PRINT_SCHEMA_FULL), + ("", _PRINT_SCHEMA_FULL), ([], _PRINT_SCHEMA_FULL), ({}, _PRINT_SCHEMA_FULL), +]) +def test_print_schema_exact_python_stdout(print_schema_service, capsys, level, expected): + spark, service = print_schema_service + df = spark.range(10) + for _ in range(2): + assert df.printSchema(level) is None + captured = capsys.readouterr() + assert captured.out == expected + assert captured.err == "" + assert service.calls >= 1 + assert service.execute_calls == 0 + + +@pytest.mark.parametrize("level", ["two", [1], object(), 1 + 2j]) +def test_print_schema_invalid_depth(print_schema_service, capsys, level): + spark, service = print_schema_service + with pytest.raises(TypeError): + spark.range(10).printSchema(level) + assert capsys.readouterr().out == "" + assert service.execute_calls == 0 + + +@pytest.mark.parametrize("failure", ["rpc", "no-result", "missing-schema", "invalid-type"]) +def test_print_schema_default_and_failure_retry(print_schema_service, capsys, failure): + spark, service = print_schema_service + df = spark.range(10) + service.failure = failure + with pytest.raises(Exception): + df.printSchema() + assert capsys.readouterr().out == "" + assert service.calls == 1 + service.failure = None + assert df.printSchema() is None + assert capsys.readouterr().out == _PRINT_SCHEMA_FULL + assert service.calls == 2 + assert service.execute_calls == 0 + + +def test_schema_tree_type_names_and_escaped_field_names(): + schema = T.StructType([ + T.StructField("line\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" + " |-- line\\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 From 5019808781d1f35c3a88db908a4517915597c57e Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 11:36:46 -0700 Subject: [PATCH 2/4] [SPARK-60695][CONNECT] Address printSchema review feedback Reuse the shared schema mock server with per-test resets and count rejected ExecutePlan calls. Render Python UDT class names in schema trees, including nested values, and preserve the generic fallback when metadata is missing. Assert repeated printing performs one schema analysis and tidy test formatting. Convert schema errors within the detached call to avoid a large-error lint. Co-authored-by: Isaac --- crates/pyspark-rs/src/dataframe.rs | 2 +- crates/spark-connect/src/types.rs | 11 + crates/spark-connect/tests/types_behavior.rs | 83 ++++++ python/tests/_mock_connect_server.py | 23 ++ python/tests/test_dropin_offline.py | 250 ++++++++----------- 5 files changed, 216 insertions(+), 153 deletions(-) diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index 540d93e2..ef536855 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1362,7 +1362,7 @@ impl PyDataFrame { /// Print the schema through Python stdout, honoring the maximum nesting depth. #[pyo3(signature = (level=None))] fn printSchema(&self, py: Python<'_>, level: Option<&Bound<'_, PyAny>>) -> PyResult<()> { - let schema = py.detach(|| self.dataframe.schema()).to_pyerr()?; + let schema = py.detach(|| self.dataframe.schema().to_pyerr())?; // Reference PySpark treats false values and nonpositive depths as unlimited. // Positive fractional depths include the partially remaining level as well. let depth = match level { diff --git a/crates/spark-connect/src/types.rs b/crates/spark-connect/src/types.rs index 5165b91d..195747a8 100644 --- a/crates/spark-connect/src/types.rs +++ b/crates/spark-connect/src/types.rs @@ -887,6 +887,17 @@ impl DataType { | 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(), } } diff --git a/crates/spark-connect/tests/types_behavior.rs b/crates/spark-connect/tests/types_behavior.rs index 0b5aff09..47c48446 100644 --- a/crates/spark-connect/tests/types_behavior.rs +++ b/crates/spark-connect/tests/types_behavior.rs @@ -321,6 +321,89 @@ fn tree_string_with_depth_truncates_nested_structs() { assert!(!shallow.contains("inner")); } +#[test] +fn tree_string_uses_python_udt_class_names() { + use spark_connect::types::StructField; + let udt = |python_class: &str| DataType::Udt { + type_str: "udt".to_string(), + jvm_class: None, + python_class: Some(python_class.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: udt("pyspark.ml.linalg.VectorUDT"), + nullable: true, + metadata: Default::default(), + }, + StructField { + name: "array".to_string(), + data_type: DataType::Array { + element_type: Box::new(udt("example.CustomType")), + 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(udt("UnqualifiedUDT")), + 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: customtype (containsNull = false)\n", + " |-- map: map (nullable = true)\n", + " | |-- key: string\n", + " | |-- value: unqualifiedudt (valueContainsNull = true)\n", + ) + ); +} + +#[test] +fn tree_string_keeps_generic_udt_name_without_python_class() { + use spark_connect::types::StructField; + for python_class in [None, Some(""), Some("example.")] { + let schema = DataType::Struct { + fields: vec![StructField { + name: "features".to_string(), + data_type: DataType::Udt { + type_str: "udt".to_string(), + jvm_class: Some("example.VectorUDT".to_string()), + python_class: python_class.map(str::to_string), + 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", + "python_class = {python_class:?}" + ); + } +} + #[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 555f829a..bf93b619 100644 --- a/python/tests/_mock_connect_server.py +++ b/python/tests/_mock_connect_server.py @@ -58,11 +58,20 @@ def serve(connection): # they wait for a response. class Service: calls = 0 + execute_calls = 0 data_type = struct_schema() failure = None block_response = False response_gate = Event() + def reset(self): + self.calls = 0 + self.execute_calls = 0 + self.data_type = struct_schema() + self.failure = None + self.block_response = False + self.response_gate.clear() + def analyze(self, request, context): self.calls += 1 if self.block_response and not self.response_gate.wait(5): @@ -81,11 +90,17 @@ def analyze(self, request, context): # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1. return message_field(2, message_field(1, self.data_type)) + def execute(self, request, context): + self.execute_calls += 1 + context.abort(grpc.StatusCode.INTERNAL, "printSchema must not execute a plan") + yield b"" + service = Service() with ThreadPoolExecutor(max_workers=4) as executor: server = grpc.server(executor) handlers = { "AnalyzePlan": grpc.unary_unary_rpc_method_handler(service.analyze), + "ExecutePlan": grpc.unary_stream_rpc_method_handler(service.execute), "ReleaseSession": grpc.unary_unary_rpc_method_handler(lambda request, context: b""), } server.add_generic_rpc_handlers( @@ -107,6 +122,9 @@ def analyze(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() @@ -130,6 +148,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 a7b92d0d..2c3db527 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -2764,13 +2764,20 @@ def test_create_dataframe_width_mismatch_is_a_normal_exception(): # --------------------------------------------------------------------------- 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 @@ -3087,138 +3094,38 @@ def allow_response_after_request(): assert service.calls == 1 -# printSchema uses a separate mock process so the real transport is exercised. -def _print_schema_varint(value): - result = bytearray() - while value > 127: - result.append((value & 127) | 128) - value >>= 7 - result.append(value) - return bytes(result) - - -def _print_schema_message(field, value): - if isinstance(value, str): - value = value.encode() - return _print_schema_varint(field * 8 + 2) + _print_schema_varint(len(value)) + value - - -def _print_schema_type(): - msg = _print_schema_message - field = _print_schema_field - integer = msg(6, b"") - long = msg(7, b"") - string = msg(11, b"") - struct = lambda *fields: msg(21, b"".join(msg(1, f) for f in fields)) - array = msg(20, msg(1, struct(field("item", long))) + b"\x10\x00") - mapping = msg(22, msg(1, string) + msg(2, struct(field("value", integer))) + b"\x18\x01") - return struct(field("outer", struct(field("inner", integer, False))), - field("items", array, False), field("lookup", mapping)) +# --------------------------------------------------------------------------- printSchema def _print_schema_field(name, data_type, nullable=True): - return (_print_schema_message(1, name) + _print_schema_message(2, data_type) + (b"\x18\x01" if nullable else b"\x18\x00") - + _print_schema_message(4, '{"note":"original"}')) - - -def _print_schema_serve(connection): - import grpc - from concurrent.futures import ThreadPoolExecutor - # A separate process also exercises columns/dtypes, whose first access may - # hold the client's GIL while waiting for the RPC. - class Service: - calls = 0 - execute_calls = 0 - data_type = _print_schema_type() - failure = None - - def analyze(self, request, context): - self.calls += 1 - if self.failure == "rpc": - context.abort(grpc.StatusCode.INVALID_ARGUMENT, "schema unavailable") - if self.failure == "no-result": - return b"" - if self.failure == "missing-schema": - return _print_schema_message(2, b"") - if self.failure == "invalid-type": - return _print_schema_message(2, _print_schema_message(1, b"")) - # base.proto: AnalyzePlanResponse.schema=2, Schema.schema=1. - # session_id is the first length-delimited field in these requests. - assert request[0] == 10 and request[1] < 128 - session_id = request[2:2 + request[1]] - return (_print_schema_message(1, session_id) - + _print_schema_message(2, _print_schema_message(1, self.data_type))) - - def execute(self, request, context): - self.execute_calls += 1 - context.abort(grpc.StatusCode.INTERNAL, "printSchema must not execute a plan") - yield b"" - - service = Service() - with ThreadPoolExecutor(max_workers=4) as executor: - server = grpc.server(executor) - handler = grpc.unary_unary_rpc_method_handler(service.analyze) - server.add_generic_rpc_handlers((grpc.method_handlers_generic_handler( - "spark.connect.SparkConnectService", {"AnalyzePlan": handler, "ExecutePlan": grpc.unary_stream_rpc_method_handler(service.execute), - "ReleaseSession": grpc.unary_unary_rpc_method_handler(lambda request, context: request[:2 + request[1]])}),)) - port = server.add_insecure_port("127.0.0.1:0") - server.start() - connection.send(port) - try: - while True: - command = connection.recv() - if command[0] == "stop": - break - if command[0] == "get": - connection.send(getattr(service, command[1])) - elif command[0] == "set": - setattr(service, command[1], command[2]) - connection.send(None) - finally: - server.stop(0).wait() - - -class _PrintSchemaService: - def __init__(self, connection): - object.__setattr__(self, "connection", connection) - - def __getattr__(self, name): - self.connection.send(("get", name)) - assert self.connection.poll(10), "Mock server did not respond" - return self.connection.recv() - - def __setattr__(self, name, value): - self.connection.send(("set", name, value)) - assert self.connection.poll(10), "Mock server did not respond" - self.connection.recv() + return ( + message_field(1, name) + + message_field(2, data_type) + + (b"\x18\x01" if nullable else b"\x18\x00") + ) -@pytest.fixture -def print_schema_service(): - import multiprocessing - pytest.importorskip("grpc") - context = multiprocessing.get_context("spawn") - parent, child = context.Pipe() - process = context.Process(target=_print_schema_serve, args=(child,)) - spark = None - process.start() - try: - assert parent.poll(20), "Mock server did not start" - port = parent.recv() - from pyspark.sql import SparkSession - spark = SparkSession.builder.remote(f"sc://127.0.0.1:{port}").create() - yield spark, _PrintSchemaService(parent) - finally: - if spark is not None: - spark.stop() - if process.is_alive(): - parent.send(("stop",)) - process.join(10) - if process.is_alive(): - process.terminate() - process.join(10) - parent.close() - child.close() +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 = ( @@ -3251,31 +3158,45 @@ def print_schema_service(): ) -@pytest.mark.parametrize("level, expected", [ - (None, _PRINT_SCHEMA_FULL), (0, _PRINT_SCHEMA_FULL), - (-1, _PRINT_SCHEMA_FULL), (-10, _PRINT_SCHEMA_FULL), - (1, _PRINT_SCHEMA_LEVEL1), (2, _PRINT_SCHEMA_LEVEL2), - (3, _PRINT_SCHEMA_FULL), (2**100, _PRINT_SCHEMA_FULL), - (False, _PRINT_SCHEMA_FULL), (True, _PRINT_SCHEMA_LEVEL1), - (0.5, _PRINT_SCHEMA_LEVEL1), (1.5, _PRINT_SCHEMA_LEVEL2), - (float("inf"), _PRINT_SCHEMA_FULL), (float("nan"), _PRINT_SCHEMA_FULL), - ("", _PRINT_SCHEMA_FULL), ([], _PRINT_SCHEMA_FULL), ({}, _PRINT_SCHEMA_FULL), -]) -def test_print_schema_exact_python_stdout(print_schema_service, capsys, level, expected): - spark, service = print_schema_service +@pytest.mark.parametrize( + "level, expected", + [ + (None, _PRINT_SCHEMA_FULL), + (0, _PRINT_SCHEMA_FULL), + (-1, _PRINT_SCHEMA_FULL), + (-10, _PRINT_SCHEMA_FULL), + (1, _PRINT_SCHEMA_LEVEL1), + (2, _PRINT_SCHEMA_LEVEL2), + (3, _PRINT_SCHEMA_FULL), + (2**100, _PRINT_SCHEMA_FULL), + (False, _PRINT_SCHEMA_FULL), + (True, _PRINT_SCHEMA_LEVEL1), + (0.5, _PRINT_SCHEMA_LEVEL1), + (1.5, _PRINT_SCHEMA_LEVEL2), + (float("inf"), _PRINT_SCHEMA_FULL), + (float("nan"), _PRINT_SCHEMA_FULL), + ("", _PRINT_SCHEMA_FULL), + ([], _PRINT_SCHEMA_FULL), + ({}, _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) for _ in range(2): assert df.printSchema(level) is None captured = capsys.readouterr() assert captured.out == expected assert captured.err == "" - assert service.calls >= 1 + assert service.calls == 1 assert service.execute_calls == 0 @pytest.mark.parametrize("level", ["two", [1], object(), 1 + 2j]) -def test_print_schema_invalid_depth(print_schema_service, capsys, level): - spark, service = print_schema_service +def test_print_schema_invalid_depth(mock_connect, capsys, level): + spark, service = mock_connect + service.data_type = _print_schema_type() with pytest.raises(TypeError): spark.range(10).printSchema(level) assert capsys.readouterr().out == "" @@ -3283,8 +3204,9 @@ def test_print_schema_invalid_depth(print_schema_service, capsys, level): @pytest.mark.parametrize("failure", ["rpc", "no-result", "missing-schema", "invalid-type"]) -def test_print_schema_default_and_failure_retry(print_schema_service, capsys, failure): - spark, service = print_schema_service +def test_print_schema_default_and_failure_retry(mock_connect, capsys, failure): + spark, service = mock_connect + service.data_type = _print_schema_type() df = spark.range(10) service.failure = failure with pytest.raises(Exception): @@ -3299,12 +3221,14 @@ def test_print_schema_default_and_failure_retry(print_schema_service, capsys, fa def test_schema_tree_type_names_and_escaped_field_names(): - schema = T.StructType([ - T.StructField("line\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()), - ]) + schema = T.StructType( + [ + T.StructField("line\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" " |-- line\\n\\r\\t\\f\\b\\v\\a: long (nullable = false)\n" @@ -3315,3 +3239,25 @@ def test_schema_tree_type_names_and_escaped_field_names(): 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"), + ("example.CustomType", "customtype"), + (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 From 559cab3ebbdfcbe5ecdc9f593c93728f2f9d8d5e Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Fri, 30 Oct 2026 12:27:07 -0700 Subject: [PATCH 3/4] [SPARK-60695][CONNECT] Avoid intermediate schema tree strings Escape field names directly into the tree buffer in one pass and write line fragments without allocating temporary formatted strings. Extend the exact output test to preserve Unicode and literal backslashes alongside controls. Co-authored-by: Isaac --- crates/spark-connect/src/types.rs | 58 ++++++++++++++++------------- python/tests/test_dropin_offline.py | 4 +- 2 files changed, 34 insertions(+), 28 deletions(-) diff --git a/crates/spark-connect/src/types.rs b/crates/spark-connect/src/types.rs index 195747a8..41679b96 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}; @@ -910,22 +910,27 @@ impl DataType { match self { DataType::Struct { fields } => { for f in fields { - let name = f - .name - .replace('\n', "\\n") - .replace('\r', "\\r") - .replace('\t', "\\t") - .replace('\x0c', "\\f") - .replace('\x08', "\\b") - .replace('\x0b', "\\v") - .replace('\x07', "\\a"); - out.push_str(&format!( - "{}-- {}: {} (nullable = {})\n", - prefix, - name, + out.push_str(prefix); + out.push_str("-- "); + for ch in f.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), + } + } + writeln!( + out, + ": {} (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); } } @@ -933,12 +938,14 @@ impl DataType { element_type, contains_null, } => { - out.push_str(&format!( - "{}-- element: {} (containsNull = {})\n", + 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 { @@ -946,18 +953,17 @@ impl DataType { value_type, value_contains_null, } => { - out.push_str(&format!( - "{}-- key: {}\n", - prefix, - key_type.tree_type_name() - )); + 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); - out.push_str(&format!( - "{}-- value: {} (valueContainsNull = {})\n", + 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); } _ => {} diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 2c3db527..51a94324 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -3223,7 +3223,7 @@ def test_print_schema_default_and_failure_retry(mock_connect, capsys, failure): def test_schema_tree_type_names_and_escaped_field_names(): schema = T.StructType( [ - T.StructField("line\n\r\t\f\b\v\a", T.LongType(), False), + 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()), @@ -3231,7 +3231,7 @@ def test_schema_tree_type_names_and_escaped_field_names(): ) expected = ( "root\n" - " |-- line\\n\\r\\t\\f\\b\\v\\a: long (nullable = false)\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" From f807ba5ceed9fe9d51f7a08722152ce0223d9834 Mon Sep 17 00:00:00 2001 From: DB Tsai Date: Sat, 31 Oct 2026 15:25:55 -0700 Subject: [PATCH 4/4] [SPARK-60695][CONNECT] Address LC schema rendering feedback Share Python depth normalization between printSchema and treeString while preserving reference defaults and argument validation. Print the schema tree for Rust DataFrame callers too. Document collection depths, extract field-name escaping, and avoid allocating prefixes for scalar leaves. Keep representative depth cases, one cache reuse check and one RPC-error propagation check. Limit UDT cases to VectorUDT and missing-class metadata. Co-authored-by: Isaac --- crates/pyspark-rs/src/dataframe.rs | 22 +++---- crates/pyspark-rs/src/types.rs | 22 ++++++- crates/spark-connect/src/dataframe.rs | 2 +- crates/spark-connect/src/types.rs | 39 +++++++----- crates/spark-connect/tests/types_behavior.rs | 53 ++++++++-------- python/tests/test_dropin_offline.py | 67 ++++++++++---------- 6 files changed, 113 insertions(+), 92 deletions(-) diff --git a/crates/pyspark-rs/src/dataframe.rs b/crates/pyspark-rs/src/dataframe.rs index ef536855..d2d40cad 100644 --- a/crates/pyspark-rs/src/dataframe.rs +++ b/crates/pyspark-rs/src/dataframe.rs @@ -1362,19 +1362,17 @@ impl PyDataFrame { /// Print the schema through Python stdout, honoring the maximum nesting depth. #[pyo3(signature = (level=None))] fn printSchema(&self, py: Python<'_>, level: Option<&Bound<'_, PyAny>>) -> PyResult<()> { - let schema = py.detach(|| self.dataframe.schema().to_pyerr())?; - // Reference PySpark treats false values and nonpositive depths as unlimited. - // Positive fractional depths include the partially remaining level as well. - let depth = match level { - Some(level) if level.is_truthy()? && level.gt(0)? => { - if level.ge(i32::MAX)? { - i32::MAX - } else { - level.extract::()?.ceil() as i32 - } - } - _ => i32::MAX, + // 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(()) diff --git a/crates/pyspark-rs/src/types.rs b/crates/pyspark-rs/src/types.rs index b58db7cb..4b9872ec 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). @@ -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 { + #[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 4ce38983..6cd18647 100644 --- a/crates/spark-connect/src/dataframe.rs +++ b/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(()) } diff --git a/crates/spark-connect/src/types.rs b/crates/spark-connect/src/types.rs index 41679b96..3dbce24f 100644 --- a/crates/spark-connect/src/types.rs +++ b/crates/spark-connect/src/types.rs @@ -858,9 +858,10 @@ 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 { .. } => { @@ -906,24 +907,13 @@ impl DataType { if remaining <= 0 { return; } - let child_prefix = format!("{prefix} |"); match self { DataType::Struct { fields } => { + let child_prefix = format!("{prefix} |"); for f in fields { out.push_str(prefix); out.push_str("-- "); - for ch in f.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), - } - } + push_escaped_name(out, &f.name); writeln!( out, ": {} (nullable = {})", @@ -938,6 +928,7 @@ impl DataType { element_type, contains_null, } => { + let child_prefix = format!("{prefix} |"); writeln!( out, "{}-- element: {} (containsNull = {})", @@ -953,6 +944,7 @@ impl DataType { 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); @@ -1043,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()) diff --git a/crates/spark-connect/tests/types_behavior.rs b/crates/spark-connect/tests/types_behavior.rs index 47c48446..56ee941f 100644 --- a/crates/spark-connect/tests/types_behavior.rs +++ b/crates/spark-connect/tests/types_behavior.rs @@ -324,10 +324,10 @@ fn tree_string_with_depth_truncates_nested_structs() { #[test] fn tree_string_uses_python_udt_class_names() { use spark_connect::types::StructField; - let udt = |python_class: &str| DataType::Udt { + let vector_udt = || DataType::Udt { type_str: "udt".to_string(), jvm_class: None, - python_class: Some(python_class.to_string()), + python_class: Some("pyspark.ml.linalg.VectorUDT".to_string()), serialized_python_class: None, sql_type: Some(Box::new(DataType::Integer)), }; @@ -335,14 +335,14 @@ fn tree_string_uses_python_udt_class_names() { fields: vec![ StructField { name: "features".to_string(), - data_type: udt("pyspark.ml.linalg.VectorUDT"), + data_type: vector_udt(), nullable: true, metadata: Default::default(), }, StructField { name: "array".to_string(), data_type: DataType::Array { - element_type: Box::new(udt("example.CustomType")), + element_type: Box::new(vector_udt()), contains_null: false, }, nullable: false, @@ -354,7 +354,7 @@ fn tree_string_uses_python_udt_class_names() { key_type: Box::new(DataType::String { collation: String::new(), }), - value_type: Box::new(udt("UnqualifiedUDT")), + value_type: Box::new(vector_udt()), value_contains_null: true, }, nullable: true, @@ -369,10 +369,10 @@ fn tree_string_uses_python_udt_class_names() { "root\n", " |-- features: vectorudt (nullable = true)\n", " |-- array: array (nullable = false)\n", - " | |-- element: customtype (containsNull = false)\n", + " | |-- element: vectorudt (containsNull = false)\n", " |-- map: map (nullable = true)\n", " | |-- key: string\n", - " | |-- value: unqualifiedudt (valueContainsNull = true)\n", + " | |-- value: vectorudt (valueContainsNull = true)\n", ) ); } @@ -380,28 +380,25 @@ fn tree_string_uses_python_udt_class_names() { #[test] fn tree_string_keeps_generic_udt_name_without_python_class() { use spark_connect::types::StructField; - for python_class in [None, Some(""), Some("example.")] { - let schema = DataType::Struct { - fields: vec![StructField { - name: "features".to_string(), - data_type: DataType::Udt { - type_str: "udt".to_string(), - jvm_class: Some("example.VectorUDT".to_string()), - python_class: python_class.map(str::to_string), - serialized_python_class: None, - sql_type: Some(Box::new(DataType::Integer)), - }, - nullable: true, - metadata: Default::default(), - }], - }; + 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", - "python_class = {python_class:?}" - ); - } + assert_eq!( + schema.tree_string().unwrap(), + "root\n |-- features: udt (nullable = true)\n", + ); } #[test] diff --git a/python/tests/test_dropin_offline.py b/python/tests/test_dropin_offline.py index 51a94324..6d5b8113 100644 --- a/python/tests/test_dropin_offline.py +++ b/python/tests/test_dropin_offline.py @@ -3161,62 +3161,66 @@ def struct(*fields): @pytest.mark.parametrize( "level, expected", [ - (None, _PRINT_SCHEMA_FULL), (0, _PRINT_SCHEMA_FULL), - (-1, _PRINT_SCHEMA_FULL), - (-10, _PRINT_SCHEMA_FULL), (1, _PRINT_SCHEMA_LEVEL1), - (2, _PRINT_SCHEMA_LEVEL2), - (3, _PRINT_SCHEMA_FULL), - (2**100, _PRINT_SCHEMA_FULL), - (False, _PRINT_SCHEMA_FULL), - (True, _PRINT_SCHEMA_LEVEL1), - (0.5, _PRINT_SCHEMA_LEVEL1), (1.5, _PRINT_SCHEMA_LEVEL2), - (float("inf"), _PRINT_SCHEMA_FULL), + (2**100, _PRINT_SCHEMA_FULL), (float("nan"), _PRINT_SCHEMA_FULL), - ("", _PRINT_SCHEMA_FULL), - ([], _PRINT_SCHEMA_FULL), - ({}, _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(level) is None - captured = capsys.readouterr() - assert captured.out == expected - assert captured.err == "" + 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", ["two", [1], object(), 1 + 2j]) -def test_print_schema_invalid_depth(mock_connect, capsys, level): +@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): - spark.range(10).printSchema(level) - assert capsys.readouterr().out == "" + df.schema.treeString(level) assert service.execute_calls == 0 -@pytest.mark.parametrize("failure", ["rpc", "no-result", "missing-schema", "invalid-type"]) -def test_print_schema_default_and_failure_retry(mock_connect, capsys, failure): +def test_schema_tree_invalid_depth(mock_connect, capsys): spark, service = mock_connect service.data_type = _print_schema_type() df = spark.range(10) - service.failure = failure - with pytest.raises(Exception): - df.printSchema() + 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.calls == 1 - service.failure = None - assert df.printSchema() is None - assert capsys.readouterr().out == _PRINT_SCHEMA_FULL - assert service.calls == 2 assert service.execute_calls == 0 @@ -3245,7 +3249,6 @@ def test_schema_tree_type_names_and_escaped_field_names(): "python_class, expected", [ ("pyspark.ml.linalg.VectorUDT", "vectorudt"), - ("example.CustomType", "customtype"), (None, "udt"), ], )