viirya commented on code in PR #131:
URL:
https://github.com/apache/spark-connect-rust/pull/131#discussion_r4431033036
##########
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 },
Review Comment:
The doc comment on `tree_string_with_depth` is stale. Truncation now applies
to array/map children as well as structs, and `max_depth <= 0` means unlimited,
where it used to print only the top level. That's a behavior change in a public
API, so it should be documented.
##########
crates/spark-connect/src/types.rs:
##########
@@ -875,33 +879,94 @@ 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;
+ }
+ let child_prefix = format!("{prefix} |");
Review Comment:
This allocates a prefix on every call, including leaf fields that return
through `_ => {}`. Can it move into the Struct/Array/Map arms?
##########
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:
Suggest stating the equivalence directly: reference does `if level:` and
then compares `> 0` at each level, so a positive fractional depth behaves as
`ceil(level)`.
##########
crates/spark-connect/src/types.rs:
##########
@@ -875,33 +879,94 @@ 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;
+ }
+ let child_prefix = format!("{prefix} |");
+ match self {
+ DataType::Struct { fields } => {
+ for f in fields {
+ out.push_str(prefix);
+ out.push_str("-- ");
+ for ch in f.name.chars() {
Review Comment:
Nit: could this be a small `push_escaped_name(out, &f.name)` helper, to
mirror `escape_meta_characters` and keep `append_tree` focused on structure?
##########
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:
Retry-after-failure is `schema()` cache behavior, and the schema-cache tests
already cover it for these four failure modes. For `printSchema`, one case that
checks the error propagates and nothing is printed should be enough.
##########
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:
Several of these cases cover the same branch: `-1`/`-10`,
`3`/`2**100`/`inf`, and `""`/`[]`/`{}`. Re-checking `calls == 1` on every case
also re-tests the #120 cache. Can we keep one case per branch and check cache
reuse once?
##########
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:
Reference can't produce an unqualified `python_class`:
`UserDefinedType.fromJson` would fail to import it. The same goes for
`Some("")` and `Some("example.")` at L383. These assertions pin behavior that
has no reference to match. Can we drop them and keep the real
`pyspark.ml.linalg.VectorUDT` and `None` cases?
--
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]
