Skip to content

Commit e55875b

Browse files
zhengruifengzhengruifeng
committed
[SPARK-48372][SPARK-45716][PYTHON] Implement StructType.treeString
### What changes were proposed in this pull request? Implement `StructType.treeString` ### Why are the changes needed? feature parity, this method is Scala-only before ### Does this PR introduce _any_ user-facing change? yes ``` In [2]: schema1 = DataType.fromDDL("c1 INT, c2 STRUCT<c3: INT, c4: STRUCT<c5: INT, c6: INT>>") In [3]: print(schema1.treeString()) root |-- c1: integer (nullable = true) |-- c2: struct (nullable = true) | |-- c3: integer (nullable = true) | |-- c4: struct (nullable = true) | | |-- c5: integer (nullable = true) | | |-- c6: integer (nullable = true) ``` ### How was this patch tested? added tests ### Was this patch authored or co-authored using generative AI tooling? no Closes #46685 from zhengruifeng/py_tree_string. Authored-by: Ruifeng Zheng <ruifengz@apache.org> Signed-off-by: Ruifeng Zheng <ruifengz@apache.org>
1 parent bf7f664 commit e55875b

3 files changed

Lines changed: 380 additions & 2 deletions

File tree

‎python/pyspark/sql/tests/test_types.py‎

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1170,6 +1170,247 @@ def test_parse_datatype_string(self):
11701170
)
11711171
self.assertEqual(VariantType(), _parse_datatype_string("variant"))
11721172

1173+
def test_tree_string(self):
1174+
schema1 = DataType.fromDDL("c1 INT, c2 STRUCT<c3: INT, c4: STRUCT<c5: INT, c6: INT>>")
1175+
1176+
self.assertEqual(
1177+
schema1.treeString().split("\n"),
1178+
[
1179+
"root",
1180+
" |-- c1: integer (nullable = true)",
1181+
" |-- c2: struct (nullable = true)",
1182+
" | |-- c3: integer (nullable = true)",
1183+
" | |-- c4: struct (nullable = true)",
1184+
" | | |-- c5: integer (nullable = true)",
1185+
" | | |-- c6: integer (nullable = true)",
1186+
"",
1187+
],
1188+
)
1189+
self.assertEqual(
1190+
schema1.treeString(-1).split("\n"),
1191+
[
1192+
"root",
1193+
" |-- c1: integer (nullable = true)",
1194+
" |-- c2: struct (nullable = true)",
1195+
" | |-- c3: integer (nullable = true)",
1196+
" | |-- c4: struct (nullable = true)",
1197+
" | | |-- c5: integer (nullable = true)",
1198+
" | | |-- c6: integer (nullable = true)",
1199+
"",
1200+
],
1201+
)
1202+
self.assertEqual(
1203+
schema1.treeString(0).split("\n"),
1204+
[
1205+
"root",
1206+
" |-- c1: integer (nullable = true)",
1207+
" |-- c2: struct (nullable = true)",
1208+
" | |-- c3: integer (nullable = true)",
1209+
" | |-- c4: struct (nullable = true)",
1210+
" | | |-- c5: integer (nullable = true)",
1211+
" | | |-- c6: integer (nullable = true)",
1212+
"",
1213+
],
1214+
)
1215+
self.assertEqual(
1216+
schema1.treeString(1).split("\n"),
1217+
[
1218+
"root",
1219+
" |-- c1: integer (nullable = true)",
1220+
" |-- c2: struct (nullable = true)",
1221+
"",
1222+
],
1223+
)
1224+
self.assertEqual(
1225+
schema1.treeString(2).split("\n"),
1226+
[
1227+
"root",
1228+
" |-- c1: integer (nullable = true)",
1229+
" |-- c2: struct (nullable = true)",
1230+
" | |-- c3: integer (nullable = true)",
1231+
" | |-- c4: struct (nullable = true)",
1232+
"",
1233+
],
1234+
)
1235+
self.assertEqual(
1236+
schema1.treeString(3).split("\n"),
1237+
[
1238+
"root",
1239+
" |-- c1: integer (nullable = true)",
1240+
" |-- c2: struct (nullable = true)",
1241+
" | |-- c3: integer (nullable = true)",
1242+
" | |-- c4: struct (nullable = true)",
1243+
" | | |-- c5: integer (nullable = true)",
1244+
" | | |-- c6: integer (nullable = true)",
1245+
"",
1246+
],
1247+
)
1248+
self.assertEqual(
1249+
schema1.treeString(4).split("\n"),
1250+
[
1251+
"root",
1252+
" |-- c1: integer (nullable = true)",
1253+
" |-- c2: struct (nullable = true)",
1254+
" | |-- c3: integer (nullable = true)",
1255+
" | |-- c4: struct (nullable = true)",
1256+
" | | |-- c5: integer (nullable = true)",
1257+
" | | |-- c6: integer (nullable = true)",
1258+
"",
1259+
],
1260+
)
1261+
1262+
schema2 = DataType.fromDDL(
1263+
"c1 INT, c2 ARRAY<STRUCT<c3: INT>>, c4 STRUCT<c5: INT, c6: ARRAY<ARRAY<INT>>>"
1264+
)
1265+
self.assertEqual(
1266+
schema2.treeString(0).split("\n"),
1267+
[
1268+
"root",
1269+
" |-- c1: integer (nullable = true)",
1270+
" |-- c2: array (nullable = true)",
1271+
" | |-- element: struct (containsNull = true)",
1272+
" | | |-- c3: integer (nullable = true)",
1273+
" |-- c4: struct (nullable = true)",
1274+
" | |-- c5: integer (nullable = true)",
1275+
" | |-- c6: array (nullable = true)",
1276+
" | | |-- element: array (containsNull = true)",
1277+
" | | | |-- element: integer (containsNull = true)",
1278+
"",
1279+
],
1280+
)
1281+
self.assertEqual(
1282+
schema2.treeString(1).split("\n"),
1283+
[
1284+
"root",
1285+
" |-- c1: integer (nullable = true)",
1286+
" |-- c2: array (nullable = true)",
1287+
" |-- c4: struct (nullable = true)",
1288+
"",
1289+
],
1290+
)
1291+
self.assertEqual(
1292+
schema2.treeString(2).split("\n"),
1293+
[
1294+
"root",
1295+
" |-- c1: integer (nullable = true)",
1296+
" |-- c2: array (nullable = true)",
1297+
" | |-- element: struct (containsNull = true)",
1298+
" |-- c4: struct (nullable = true)",
1299+
" | |-- c5: integer (nullable = true)",
1300+
" | |-- c6: array (nullable = true)",
1301+
"",
1302+
],
1303+
)
1304+
self.assertEqual(
1305+
schema2.treeString(3).split("\n"),
1306+
[
1307+
"root",
1308+
" |-- c1: integer (nullable = true)",
1309+
" |-- c2: array (nullable = true)",
1310+
" | |-- element: struct (containsNull = true)",
1311+
" | | |-- c3: integer (nullable = true)",
1312+
" |-- c4: struct (nullable = true)",
1313+
" | |-- c5: integer (nullable = true)",
1314+
" | |-- c6: array (nullable = true)",
1315+
" | | |-- element: array (containsNull = true)",
1316+
"",
1317+
],
1318+
)
1319+
self.assertEqual(
1320+
schema2.treeString(4).split("\n"),
1321+
[
1322+
"root",
1323+
" |-- c1: integer (nullable = true)",
1324+
" |-- c2: array (nullable = true)",
1325+
" | |-- element: struct (containsNull = true)",
1326+
" | | |-- c3: integer (nullable = true)",
1327+
" |-- c4: struct (nullable = true)",
1328+
" | |-- c5: integer (nullable = true)",
1329+
" | |-- c6: array (nullable = true)",
1330+
" | | |-- element: array (containsNull = true)",
1331+
" | | | |-- element: integer (containsNull = true)",
1332+
"",
1333+
],
1334+
)
1335+
1336+
schema3 = DataType.fromDDL(
1337+
"c1 MAP<INT, STRUCT<c2: MAP<INT, INT>>>, c3 STRUCT<c4: MAP<INT, MAP<INT, INT>>>"
1338+
)
1339+
self.assertEqual(
1340+
schema3.treeString(0).split("\n"),
1341+
[
1342+
"root",
1343+
" |-- c1: map (nullable = true)",
1344+
" | |-- key: integer",
1345+
" | |-- value: struct (valueContainsNull = true)",
1346+
" | | |-- c2: map (nullable = true)",
1347+
" | | | |-- key: integer",
1348+
" | | | |-- value: integer (valueContainsNull = true)",
1349+
" |-- c3: struct (nullable = true)",
1350+
" | |-- c4: map (nullable = true)",
1351+
" | | |-- key: integer",
1352+
" | | |-- value: map (valueContainsNull = true)",
1353+
" | | | |-- key: integer",
1354+
" | | | |-- value: integer (valueContainsNull = true)",
1355+
"",
1356+
],
1357+
)
1358+
self.assertEqual(
1359+
schema3.treeString(1).split("\n"),
1360+
[
1361+
"root",
1362+
" |-- c1: map (nullable = true)",
1363+
" |-- c3: struct (nullable = true)",
1364+
"",
1365+
],
1366+
)
1367+
self.assertEqual(
1368+
schema3.treeString(2).split("\n"),
1369+
[
1370+
"root",
1371+
" |-- c1: map (nullable = true)",
1372+
" | |-- key: integer",
1373+
" | |-- value: struct (valueContainsNull = true)",
1374+
" |-- c3: struct (nullable = true)",
1375+
" | |-- c4: map (nullable = true)",
1376+
"",
1377+
],
1378+
)
1379+
self.assertEqual(
1380+
schema3.treeString(3).split("\n"),
1381+
[
1382+
"root",
1383+
" |-- c1: map (nullable = true)",
1384+
" | |-- key: integer",
1385+
" | |-- value: struct (valueContainsNull = true)",
1386+
" | | |-- c2: map (nullable = true)",
1387+
" |-- c3: struct (nullable = true)",
1388+
" | |-- c4: map (nullable = true)",
1389+
" | | |-- key: integer",
1390+
" | | |-- value: map (valueContainsNull = true)",
1391+
"",
1392+
],
1393+
)
1394+
self.assertEqual(
1395+
schema3.treeString(4).split("\n"),
1396+
[
1397+
"root",
1398+
" |-- c1: map (nullable = true)",
1399+
" | |-- key: integer",
1400+
" | |-- value: struct (valueContainsNull = true)",
1401+
" | | |-- c2: map (nullable = true)",
1402+
" | | | |-- key: integer",
1403+
" | | | |-- value: integer (valueContainsNull = true)",
1404+
" |-- c3: struct (nullable = true)",
1405+
" | |-- c4: map (nullable = true)",
1406+
" | | |-- key: integer",
1407+
" | | |-- value: map (valueContainsNull = true)",
1408+
" | | | |-- key: integer",
1409+
" | | | |-- value: integer (valueContainsNull = true)",
1410+
"",
1411+
],
1412+
)
1413+
11731414
def test_metadata_null(self):
11741415
schema = StructType(
11751416
[

‎python/pyspark/sql/types.py‎

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,12 @@
4747

4848
from pyspark.util import is_remote_only
4949
from pyspark.serializers import CloudPickleSerializer
50-
from pyspark.sql.utils import has_numpy, get_active_spark_context
50+
from pyspark.sql.utils import (
51+
has_numpy,
52+
get_active_spark_context,
53+
escape_meta_characters,
54+
StringConcat,
55+
)
5156
from pyspark.sql.variant_utils import VariantUtils
5257
from pyspark.errors import (
5358
PySparkNotImplementedError,
@@ -99,6 +104,8 @@
99104
"VariantVal",
100105
]
101106

107+
_JVM_INT_MAX: int = (1 << 31) - 1
108+
102109

103110
class DataType:
104111
"""Base class for data types."""
@@ -199,6 +206,17 @@ def fromDDL(cls, ddl: str) -> "DataType":
199206
assert len(schema) == 1
200207
return schema[0].dataType
201208

209+
@classmethod
210+
def _data_type_build_formatted_string(
211+
cls,
212+
dataType: "DataType",
213+
prefix: str,
214+
stringConcat: StringConcat,
215+
maxDepth: int,
216+
) -> None:
217+
if isinstance(dataType, (ArrayType, StructType, MapType)):
218+
dataType._build_formatted_string(prefix, stringConcat, maxDepth - 1)
219+
202220

203221
# This singleton pattern does not work with pickle, you will get
204222
# another object after pickle and unpickle
@@ -734,6 +752,21 @@ def fromInternal(self, obj: List[Optional[T]]) -> List[Optional[T]]:
734752
return obj
735753
return obj and [self.elementType.fromInternal(v) for v in obj]
736754

755+
def _build_formatted_string(
756+
self,
757+
prefix: str,
758+
stringConcat: StringConcat,
759+
maxDepth: int = _JVM_INT_MAX,
760+
) -> None:
761+
if maxDepth > 0:
762+
stringConcat.append(
763+
f"{prefix}-- element: {self.elementType.typeName()} "
764+
+ f"(containsNull = {str(self.containsNull).lower()})\n"
765+
)
766+
DataType._data_type_build_formatted_string(
767+
self.elementType, f"{prefix} |", stringConcat, maxDepth
768+
)
769+
737770

738771
class MapType(DataType):
739772
"""Map data type.
@@ -868,6 +901,25 @@ def fromInternal(self, obj: Dict[T, Optional[U]]) -> Dict[T, Optional[U]]:
868901
(self.keyType.fromInternal(k), self.valueType.fromInternal(v)) for k, v in obj.items()
869902
)
870903

904+
def _build_formatted_string(
905+
self,
906+
prefix: str,
907+
stringConcat: StringConcat,
908+
maxDepth: int = _JVM_INT_MAX,
909+
) -> None:
910+
if maxDepth > 0:
911+
stringConcat.append(f"{prefix}-- key: {self.keyType.typeName()}\n")
912+
DataType._data_type_build_formatted_string(
913+
self.keyType, f"{prefix} |", stringConcat, maxDepth
914+
)
915+
stringConcat.append(
916+
f"{prefix}-- value: {self.valueType.typeName()} "
917+
+ f"(valueContainsNull = {str(self.valueContainsNull).lower()})\n"
918+
)
919+
DataType._data_type_build_formatted_string(
920+
self.valueType, f"{prefix} |", stringConcat, maxDepth
921+
)
922+
871923

872924
class StructField(DataType):
873925
"""A field in :class:`StructType`.
@@ -1016,6 +1068,21 @@ def typeName(self) -> str: # type: ignore[override]
10161068
message_parameters={},
10171069
)
10181070

1071+
def _build_formatted_string(
1072+
self,
1073+
prefix: str,
1074+
stringConcat: StringConcat,
1075+
maxDepth: int = _JVM_INT_MAX,
1076+
) -> None:
1077+
if maxDepth > 0:
1078+
stringConcat.append(
1079+
f"{prefix}-- {escape_meta_characters(self.name)}: {self.dataType.typeName()} "
1080+
+ f"(nullable = {str(self.nullable).lower()})\n"
1081+
)
1082+
DataType._data_type_build_formatted_string(
1083+
self.dataType, f"{prefix} |", stringConcat, maxDepth
1084+
)
1085+
10191086

10201087
class StructType(DataType):
10211088
"""Struct type, consisting of a list of :class:`StructField`.
@@ -1436,6 +1503,24 @@ def fromInternal(self, obj: Tuple) -> "Row":
14361503
values = obj
14371504
return _create_row(self.names, values)
14381505

1506+
def _build_formatted_string(
1507+
self,
1508+
prefix: str,
1509+
stringConcat: StringConcat,
1510+
maxDepth: int = _JVM_INT_MAX,
1511+
) -> None:
1512+
for field in self.fields:
1513+
field._build_formatted_string(prefix, stringConcat, maxDepth)
1514+
1515+
def treeString(self, maxDepth: int = _JVM_INT_MAX) -> str:
1516+
stringConcat = StringConcat()
1517+
stringConcat.append("root\n")
1518+
prefix = " |"
1519+
depth = maxDepth if maxDepth > 0 else _JVM_INT_MAX
1520+
for field in self.fields:
1521+
field._build_formatted_string(prefix, stringConcat, depth)
1522+
return stringConcat.toString()
1523+
14391524

14401525
class VariantType(AtomicType):
14411526
"""

0 commit comments

Comments
 (0)