dbtsai commented on code in PR #131:
URL:
https://github.com/apache/spark-connect-rust/pull/131#discussion_r4433528788
##########
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<i32>) -> 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.
Review Comment:
Stated the ceil(level) equivalence directly and moved normalization into a
helper shared by printSchema and StructType.treeString. The printSchema
truthiness check stays separate to preserve reference defaults.
##########
python/tests/test_dropin_offline.py:
##########
@@ -3085,3 +3092,172 @@ def allow_response_after_request():
else:
assert result == [("id", "bigint")]
assert service.calls == 1
+
+
+# ---------------------------------------------------------------------------
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"
+)
+
+
[email protected](
+ "level, expected",
+ [
+ (None, _PRINT_SCHEMA_FULL),
+ (0, _PRINT_SCHEMA_FULL),
+ (-1, _PRINT_SCHEMA_FULL),
+ (-10, _PRINT_SCHEMA_FULL),
Review Comment:
Reduced the matrix to representative depth cases, exercising both
printSchema and treeString against the same exact output. Cache reuse is
checked once in a separate default-depth test.
##########
python/tests/test_dropin_offline.py:
##########
@@ -3085,3 +3092,172 @@ def allow_response_after_request():
else:
assert result == [("id", "bigint")]
assert service.calls == 1
+
+
+# ---------------------------------------------------------------------------
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"
+)
+
+
[email protected](
+ "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.execute_calls == 0
+
+
[email protected]("level", ["two", [1], object(), 1 + 2j])
+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 == ""
+ assert service.execute_calls == 0
+
+
[email protected]("failure", ["rpc", "no-result", "missing-schema",
"invalid-type"])
Review Comment:
Replaced the four retry cases with one RPC-error propagation test that
checks no output is printed. The existing schema-cache tests retain retry
coverage.
##########
crates/spark-connect/tests/types_behavior.rs:
##########
@@ -314,13 +314,96 @@ 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 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")),
Review Comment:
Removed the unqualified, empty and trailing-dot class cases and the
fabricated custom class. UDT tests now use pyspark.ml.linalg.VectorUDT directly
and inside arrays/maps, plus the missing-Python-class fallback.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
