Support Iceberg in read_from/write_from JobSocket (#147)

Author: george-zubrienko

spark_utils/common/functions.py

@@ -76,7 +76,7 @@ def read_from_socket(
     :return: Spark dataframe
     """
     read_options = read_options or {}
-    if socket.data_format.startswith("hive"):
+    if socket.data_format.startswith("hive") or socket.data_format.startswith("iceberg"):
         return spark_session.table(socket.data_path)
 
     return spark_session.read.options(**read_options).format(socket.data_format).load(socket.data_path)
@@ -103,6 +103,11 @@ def write_to_socket(
     if partition_count:
         data = data.repartition(partition_count, *partition_by)
 
+    # ignore all external write options as Iceberg writer will take care of those
+    if socket.data_format.startswith("iceberg"):
+        data.writeTo(socket.data_path).createOrReplace()
+        return
+
     writer = data.write.mode("overwrite").options(**write_options)
 
     if partition_by:

test/test_iceberg.py

@@ -1,6 +1,9 @@
-import pytest
 from pyspark.sql import SparkSession, DataFrame
 
+from spark_utils.common.functions import write_to_socket, read_from_socket
+from spark_utils.models.job_socket import JobSocket
+from test.test_common_functions import are_dfs_equal
+
 
 def test_iceberg_rest_create_schema(iceberg_spark_session: SparkSession):
     try:
@@ -19,3 +22,28 @@ def test_iceberg_rest_create_table(iceberg_spark_session: SparkSession):
         assert rows.collect()[0].asDict() == {"C0": 1, "C1": "1231", "C2": 1.0}
     except BaseException as e:
         raise RuntimeError("Failed to create table") from e
+
+
+def test_write_to_socket(
+    iceberg_spark_session: SparkSession,
+):
+    output_socket = JobSocket(
+        alias="test",
+        data_path=f"iceberg.test.job_socket_write",
+        data_format="iceberg",
+    )
+    df = iceberg_spark_session.createDataFrame(
+        [{"C0": 1, "C1": "1231", "C2": 1.0}, {"C0": 2, "C1": "1232", "C2": 2.0}, {"C0": 3, "C1": "1233", "C2": 3.0}]
+    )
+
+    write_to_socket(
+        data=df,
+        socket=output_socket,
+        write_options=None,
+        partition_by=None,
+        partition_count=None,
+    )
+
+    df_read = read_from_socket(socket=output_socket, spark_session=iceberg_spark_session, read_options=None)
+
+    assert are_dfs_equal(df, df_read.select(df.columns))