Package: pbspark
pbsparkThis package provides a way to convert protobuf messages into pyspark dataframes and vice versa using pysparkudfs.InstallationTo install:pipinstallpbsparkUsageSuppose we have a pyspark DataFrame which contains a columnvaluewhich has protobuf encoded messages of ourSimpleMessage:syntax="proto3";packageexample;messageSimpleMessage{stringname=1;int64quantity=2;floatmeasure=3;}Basic conversion functionsThere are two functions for operating on columns,to_protobufandfrom_protobuf. These operations convert to/from an encoded protobuf column to a column of a struct representing the inferred message structure.MessageConverterinstances (discussed below) can optionally be passed to these functions.frompyspark.sql.sessionimportSparkSessionfromexample.example_pb2importSimpleMessagefrompbsparkimportfrom_protobuffrompbsparkimportto_protobufspark=SparkSession.builder.getOrCreate()example=SimpleMessage(name="hello",quantity=5,measure=12.3)data=[{"value":example.SerializeToString()}]df_encoded=spark.createDataFrame(data)df_decoded=df_encoded.select(from_protobuf(df_encoded.value,SimpleMessage).alias("value"))df_expanded=df_decoded.select("value.*")df_expanded.show()# +-----+--------+-------+# | name|quantity|measure|# +-----+--------+-------+# |hello|       5|   12.3|# +-----+--------+-------+df_reencoded=df_decoded.select(to_protobuf(df_decoded.value,SimpleMessage).alias("value"))There are two helper functions,df_to_protobufanddf_from_protobuffor use on dataframes. They have a kwargexpanded, which will also take care of expanding/contracting the data between the singlevaluecolumn used in these examples and a dataframe which contains a column for each message field.MessageConverterinstances (discussed below) can optionally be passed to these functions.frompyspark.sql.sessionimportSparkSessionfromexample.example_pb2importSimpleMessagefrompbsparkimportdf_from_protobuffrompbsparkimportdf_to_protobufspark=SparkSession.builder.getOrCreate()example=SimpleMessage(name="hello",quantity=5,measure=12.3)data=[{"value":example.SerializeToString()}]df_encoded=spark.createDataFrame(data)# expanded=True will perform a `.select("value.*")` after converting,# resulting in each protobuf field having its own columndf_expanded=df_from_protobuf(df_encoded,SimpleMessage,expanded=True)df_expanded.show()# +-----+--------+-------+# | name|quantity|measure|# +-----+--------+-------+# |hello|       5|   12.3|# +-----+--------+-------+# expanded=True will first pack data using `struct([df[c] for c in df.columns])`,# use this if the passed dataframe is already expandeddf_reencoded=df_to_protobuf(df_expanded,SimpleMessage,expanded=True)Column conversion using theMessageConverterThe four helper functions above are also available as methods on theMessageConverterclass. Using an instance ofMessageConverterwe can decode the column of encoded messages into a column of sparkStructTypeand then expand the fields.frompyspark.sql.sessionimportSparkSessionfrompbsparkimportMessageConverterfromexample.example_pb2importSimpleMessagespark=SparkSession.builder.getOrCreate()example=SimpleMessage(name="hello",quantity=5,measure=12.3)data=[{"value":example.SerializeToString()}]df_encoded=spark.createDataFrame(data)mc=MessageConverter()df_decoded=df_encoded.select(mc.from_protobuf(df_encoded.value,SimpleMessage).alias("value"))df_expanded=df_decoded.select("value.*")df_expanded.show()# +-----+--------+-------+# | name|quantity|measure|# +-----+--------+-------+# |hello|       5|   12.3|# +-----+--------+-------+df_expanded.schema# StructType(List(StructField(name,StringType,true),StructField(quantity,IntegerType,true),StructField(measure,FloatType,true))We can also re-encode them into protobuf.df_reencoded=df_decoded.select(mc.to_protobuf(df_decoded.value,SimpleMessage).alias("value"))For expanded data, we can also encode after packing into a struct column:frompyspark.sql.functionsimportstructdf_unexpanded=df_expanded.select(struct([df_expanded[c]forcindf_expanded.columns]).alias("value"))df_reencoded=df_unexpanded.select(mc.to_protobuf(df_unexpanded.value,SimpleMessage).alias("value"))Conversion detailsInternally,pbsparkuses protobuf'sMessageToDict, which deserializes everything into JSON compatible objects by default. The exceptions areprotobuf's bytes type, whichMessageToDictwould decode to a base64-encoded string;pbsparkwill decode any bytes fields directly to a sparkBinaryType.protobuf's well known type, Timestamp type, whichMessageToDictwould decode to a string;pbsparkwill decode any Timestamp messages directly to a sparkTimestampType(via python datetime objects).protobuf's int64 types, whichMessageToDictwould decode to a string for compatibility reasons;pbsparkwill decode these toLongType.Custom conversion of message typesCustom serde is also supported. Suppose we use ourNestedMessagefrom the repository's example and we want to serialize the key and value together into a single string.messageNestedMessage{stringkey=1;stringvalue=2;}We can create and register a custom serializer with theMessageConverter.frompbsparkimportMessageConverterfromexample.example_pb2importExampleMessagefromexample.example_pb2importNestedMessagefrompyspark.sql.typesimportStringTypemc=MessageConverter()# register a custom serializer# this will serialize the NestedMessages into a string rather than a# struct with `key` and `value` fieldsencode_nested=lambdamessage:message.key+":"+message.valuemc.register_serializer(NestedMessage,encode_nested,StringType())# ...frompyspark.sql.sessionimportSparkSessionfrompysparkimportSparkContextfrompyspark.serializersimportCloudPickleSerializersc=SparkContext(serializer=CloudPickleSerializer())spark=SparkSession(sc).builder.getOrCreate()message=ExampleMessage(nested=NestedMessage(key="hello",value="world"))data=[{"value":message.SerializeToString()}]df_encoded=spark.createDataFrame(data)df_decoded=df_encoded.select(mc.from_protobuf(df_encoded.value,ExampleMessage).alias("value"))# rather than a struct the value of `nested` is a stringdf_decoded.select("value.nested").show()# +-----------+# |     nested|# +-----------+# |hello:world|# +-----------+How to write conversion functionsMore generally, custom serde functions should be written in the following format.# Encoding takes a message instance and returns the result# of the custom transformation.defencode_nested(message:NestedMessage)->str:returnmessage.key+":"+message.value# Decoding takes the encoded value, a message instance, and path string# and populates the fields of the message instance. It returns `None`.# The path str is used in the protobuf parser to log parse error info.# Note that the first argument type should match the return type of the# encoder if using both.defdecode_nested(s:str,message:NestedMessage,path:str):key,value=s.split(":")message.key=keymessage.value=valueAvoiding PicklingErrorsA seemingly common issue with protobuf and distributed processing is when aPicklingErroris encountered when transmitting (pickling) protobuf message types from a main process to a fork. To avoid this, you need to ensure that the fully qualified module name in your protoc-generated python file is the same as the module path from which the message type is imported. In other words, for the example here, the descriptor module passed to the builder isexample.example_pb2# from example/example_pb2.py_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR,"example.example_pb2",globals())^^^^^^^^^^^^^^^^^^^And to import the message type we would call the same module path:fromexample.example_pb2importExampleMessage^^^^^^^^^^^^^^^^^^^Note that the import module is the same as the one passed to the builder from the protoc-generated python. If these do not match, then you will encounter aPicklingError. From the pickle documentation:pickle can save and restore class instances transparently, however the class definition must be importable and live in the same module as when the object was stored.To ensure that the module path is correct, you should runprotocfrom the relative root path of your proto files. For example, in this project, in theMakefileunder thegencommand, we callprotocfrom the project root rather than from within theexampledirectory.exportPROTO_PATH=.gen:poetryrunprotoc-I$$PROTO_PATH--python_out=$$PROTO_PATH--mypy_out=$$PROTO_PATH--proto_path=$$PROTO_PATH$$PROTO_PATH/example/*.protoKnown issuesRecursionErrorwhen using self-referencing protobuf messages. Spark schemas do not allow for arbitrary depth, so protobuf messages which are circular- or self-referencing will result in infinite recursion errors when inferring the schema. If you have message structures like this you should resort to creating custom conversion functions, which forcibly limit the structural depth when converting these messages.DevelopmentEnsure thatasdfis installed, then runmake setup.To format codemake fmtTo test codemake testTo run protocmake gen