[SPARK-42437][CONNECT][PYTHON][FOLLOW-UP] Storage level proto converters

### What changes were proposed in this pull request?
Converters between Proto and StorageLevel to avoid code duplication
It's follow up from apache#40015

### Why are the changes needed?
Code deduplication

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?
Existing tests

Closes apache#40859 from khalidmammadov/storage_level_converter.

Authored-by: khalidmammadov <[email protected]>
Signed-off-by: Takuya UESHIN <[email protected]>

python/pyspark/sql/connect/client.py

@@ -60,6 +60,7 @@
 from google.rpc import error_details_pb2
 
 from pyspark.resource.information import ResourceInformation
+from pyspark.sql.connect.conversion import storage_level_to_proto, proto_to_storage_level
 import pyspark.sql.connect.proto as pb2
 import pyspark.sql.connect.proto.base_pb2_grpc as grpc_lib
 import pyspark.sql.connect.types as types
@@ -469,13 +470,7 @@ def fromProto(cls, pb: Any) -> "AnalyzeResult":
         elif pb.HasField("unpersist"):
             pass
         elif pb.HasField("get_storage_level"):
-            storage_level = StorageLevel(
-                useDisk=pb.get_storage_level.storage_level.use_disk,
-                useMemory=pb.get_storage_level.storage_level.use_memory,
-                useOffHeap=pb.get_storage_level.storage_level.use_off_heap,
-                deserialized=pb.get_storage_level.storage_level.deserialized,
-                replication=pb.get_storage_level.storage_level.replication,
-            )
+            storage_level = proto_to_storage_level(pb.get_storage_level.storage_level)
         else:
             raise SparkConnectException("No analyze result found!")
 
@@ -877,15 +872,7 @@ def _analyze(self, method: str, **kwargs: Any) -> AnalyzeResult:
             req.persist.relation.CopyFrom(cast(pb2.Relation, kwargs.get("relation")))
             if kwargs.get("storage_level", None) is not None:
                 storage_level = cast(StorageLevel, kwargs.get("storage_level"))
-                req.persist.storage_level.CopyFrom(
-                    pb2.StorageLevel(
-                        use_disk=storage_level.useDisk,
-                        use_memory=storage_level.useMemory,
-                        use_off_heap=storage_level.useOffHeap,
-                        deserialized=storage_level.deserialized,
-                        replication=storage_level.replication,
-                    )
-                )
+                req.persist.storage_level.CopyFrom(storage_level_to_proto(storage_level))
         elif method == "unpersist":
             req.unpersist.relation.CopyFrom(cast(pb2.Relation, kwargs.get("relation")))
             if kwargs.get("blocking", None) is not None:

python/pyspark/sql/connect/conversion.py

@@ -43,7 +43,9 @@
     cast,
 )
 
+from pyspark.storagelevel import StorageLevel
 from pyspark.sql.connect.types import to_arrow_schema
+import pyspark.sql.connect.proto as pb2
 
 from typing import (
     Any,
@@ -486,3 +488,25 @@ def convert(table: "pa.Table", schema: StructType) -> List[Row]:
             values = [field_converters[j](columnar_data[j][i]) for j in range(table.num_columns)]
             rows.append(_create_row(fields=schema.fieldNames(), values=values))
         return rows
+
+
+def storage_level_to_proto(storage_level: StorageLevel) -> pb2.StorageLevel:
+    assert storage_level is not None and isinstance(storage_level, StorageLevel)
+    return pb2.StorageLevel(
+        use_disk=storage_level.useDisk,
+        use_memory=storage_level.useMemory,
+        use_off_heap=storage_level.useOffHeap,
+        deserialized=storage_level.deserialized,
+        replication=storage_level.replication,
+    )
+
+
+def proto_to_storage_level(storage_level: pb2.StorageLevel) -> StorageLevel:
+    assert storage_level is not None and isinstance(storage_level, pb2.StorageLevel)
+    return StorageLevel(
+        useDisk=storage_level.use_disk,
+        useMemory=storage_level.use_memory,
+        useOffHeap=storage_level.use_off_heap,
+        deserialized=storage_level.deserialized,
+        replication=storage_level.replication,
+    )

python/pyspark/sql/connect/plan.py

@@ -30,6 +30,7 @@
 from pyspark.sql.types import DataType
 
 import pyspark.sql.connect.proto as proto
+from pyspark.sql.connect.conversion import storage_level_to_proto
 from pyspark.sql.connect.column import Column
 from pyspark.sql.connect.expressions import (
     SortOrder,
@@ -1896,15 +1897,7 @@ def __init__(self, table_name: str, storage_level: Optional[StorageLevel] = None
     def plan(self, session: "SparkConnectClient") -> proto.Relation:
         _cache_table = proto.CacheTable(table_name=self._table_name)
         if self._storage_level:
-            _cache_table.storage_level.CopyFrom(
-                proto.StorageLevel(
-                    use_disk=self._storage_level.useDisk,
-                    use_memory=self._storage_level.useMemory,
-                    use_off_heap=self._storage_level.useOffHeap,
-                    deserialized=self._storage_level.deserialized,
-                    replication=self._storage_level.replication,
-                )
-            )
+            _cache_table.storage_level.CopyFrom(storage_level_to_proto(self._storage_level))
         plan = proto.Relation(catalog=proto.Catalog(cache_table=_cache_table))
         return plan