Updated sink writer with formatting changes

morpheus/stages/output/write_to_deltalake_stage.py

@@ -14,32 +14,26 @@
 
 import logging
 import os
-from mrc.core import operators as ops
+import typing
 
 import mrc
-import cudf
-import typing
+from mrc.core import operators as ops
+from pyspark.sql import SparkSession
+from pyspark.sql.types import BooleanType
+from pyspark.sql.types import DoubleType
+from pyspark.sql.types import FloatType
+from pyspark.sql.types import IntegerType
+from pyspark.sql.types import LongType
+from pyspark.sql.types import StringType
+from pyspark.sql.types import StructField
+from pyspark.sql.types import StructType
+from pyspark.sql.types import TimestampType
+
 from morpheus.cli.register_stage import register_stage
 from morpheus.config import Config
+from morpheus.messages import MessageMeta
 from morpheus.pipeline.single_port_stage import SinglePortStage
 from morpheus.pipeline.stream_pair import StreamPair
-from morpheus.messages import MessageMeta
-from morpheus.io import serializers
-
-
-from pyspark.sql import SparkSession
-from pyspark.sql.types import (
-    BooleanType,
-    DoubleType,
-    FloatType,
-    IntegerType,
-    LongType,
-    StringType,
-    StructField,
-    StructType,
-    TimestampType,
-)
-
 
 logger = logging.getLogger(__name__)
 
@@ -78,15 +72,18 @@ def __init__(self,
 
         super().__init__(config)
         self.delta_path = delta_path
-        self._configure_databricks_connect(databricks_host, databricks_token, databricks_cluster_id, databricks_port,databricks_org_id)
+        self._configure_databricks_connect(databricks_host,
+                                           databricks_token,
+                                           databricks_cluster_id,
+                                           databricks_port,
+                                           databricks_org_id)
         self.spark = SparkSession.builder.config("spark.databricks.delta.optimizeWrite.enabled", "true")\
-        .config("spark.databricks.delta.autoCompact.enabled", "true")\
-        .getOrCreate()
+            .config("spark.databricks.delta.autoCompact.enabled", "true")\
+            .getOrCreate()
 
         # Enable Arrow-based columnar data transfers
         self.spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
 
-
     @property
     def name(self) -> str:
         return "to-delta"
@@ -106,16 +103,19 @@ def supports_cpp_node(self) -> bool:
 
     def _build_single(self, builder: mrc.Builder, input_stream: StreamPair) -> StreamPair:
         stream = input_stream[0]
+
         def node_fn(obs: mrc.Observable, sub: mrc.Subscriber):
+
             def write_to_deltalake(x: MessageMeta):
                 # convert cudf to spark dataframe
                 df = x.df.to_pandas()
                 schema = self._extract_schema_from_pandas_dataframe(df)
-                spark_df = self.spark.createDataFrame(df,schema=schema)
+                spark_df = self.spark.createDataFrame(df, schema=schema)
                 spark_df.write.format('delta')\
-                .option("mergeSchema", "true").mode("append")\
-                .save(self.delta_path)
+                    .option("mergeSchema", "true").mode("append")\
+                    .save(self.delta_path)
                 return x
+
             obs.pipe(ops.map(write_to_deltalake)).subscribe(sub)
 
         to_delta = builder.make_node_full(self.unique_name, node_fn)
@@ -152,28 +152,32 @@ def _extract_schema_from_pandas_dataframe(self, df):
             spark_schema.append(spark_dtype)
         return StructType(spark_schema)
 
-
-    def _configure_databricks_connect(self,databricks_host, databricks_token, databricks_cluster_id, databricks_port, databricks_org_id):
-        if(os.environ.get('DATABRICKS_HOST',None)==None and databricks_host==None):
+    def _configure_databricks_connect(self,
+                                      databricks_host,
+                                      databricks_token,
+                                      databricks_cluster_id,
+                                      databricks_port,
+                                      databricks_org_id):
+        if (os.environ.get('DATABRICKS_HOST', None) is None and databricks_host is None):
             raise Exception("Parameter for databricks host not provided")
-        if(os.environ.get('DATABRICKS_TOKEN',None)==None and databricks_token==None):
+        if (os.environ.get('DATABRICKS_TOKEN', None) is None and databricks_token is None):
             raise Exception("Parameter for databricks token not provided")
-        if(os.environ.get('DATABRICKS_CLUSTER_ID',None)==None and databricks_cluster_id==None):
+        if (os.environ.get('DATABRICKS_CLUSTER_ID', None) is None and databricks_cluster_id is None):
             raise Exception("Parameter for databricks cluster not provided")
         host = None
         cluster = None
         token = None
         config_file = "/root/.databricks-connect"
         should_add = False
-        if(os.environ.get('DATABRICKS_HOST',None)!=None):
+        if (os.environ.get('DATABRICKS_HOST', None) is not None):
             host = os.environ.get('DATABRICKS_HOST')
         else:
             host = databricks_host
-        if(os.environ.get('DATABRICKS_TOKEN',None)!=None):
+        if (os.environ.get('DATABRICKS_TOKEN', None) is not None):
             token = os.environ.get('DATABRICKS_TOKEN')
         else:
             token = databricks_token
-        if(os.environ.get('DATABRICKS_CLUSTER_ID',None)!=None):
+        if (os.environ.get('DATABRICKS_CLUSTER_ID', None) is not None):
             cluster = os.environ.get('DATABRICKS_CLUSTER_ID')
         else:
             cluster = databricks_cluster_id
@@ -184,7 +188,8 @@ def _configure_databricks_connect(self,databricks_host, databricks_token, databr
                       "org_id": "@org_id",
                       "port": "@port"
                 }"""
-        config = config.replace("@host",host).replace("@token",token).replace("@cluster_id",cluster).replace("@org_id",databricks_org_id).replace("@port",databricks_port)
+        config = config.replace("@host", host).replace("@token", token).replace("@cluster_id", cluster).replace(
+            "@org_id", databricks_org_id).replace("@port", databricks_port)
 
         # check if the config file for databricks connect already exists
         config_exist = os.path.exists(config_file)
@@ -199,6 +204,6 @@ def _configure_databricks_connect(self,databricks_host, databricks_token, databr
         else:
             should_add = True
         if should_add:
-            with open(config_file,"w+") as f:
+            with open(config_file, "w+") as f:
                 f.write(config)
             logger.info("Databricks-connect successfully configured!")