[SPARK-40655][PYTHON][PROTOBUF] PySpark support for from_protobuf and to_protobuf

[SandishKumarHN]

From SandishKumarHN([email protected]) and Mohan Parthasarathy([email protected])

This PR follows main PR #37972

The following is an example of how to use from_protobuf and to_protobuf in Pyspark.

data = [("1", (2, "Alice", 109200))]
ddl_schema = "key STRING, value STRUCT<age: INTEGER, name: STRING, score: LONG>"
df = spark.createDataFrame(data, ddl_schema)
desc_hex = str('0ACE010A41636F6E6E6563746F722F70726F746F6275662F7372632F746573742F726'
... '5736F75726365732F70726F746F6275662F7079737061726B5F746573742E70726F746F121D6F72672E61'
... '70616368652E737061726B2E73716C2E70726F746F627566224B0A0D53696D706C654D657373616765121'
... '00A03616765180120012805520361676512120A046E616D6518022001280952046E616D6512140A057363'
... '6F7265180320012803520573636F72654215421353696D706C654D65737361676550726F746F736206707'
... '26F746F33')
 import tempfile
# Writing a protobuf description into a file, generated by using
# connector/protobuf/src/test/resources/protobuf/pyspark_test.proto file
with tempfile.TemporaryDirectory() as tmp_dir:
...     desc_file_path = "%s/pyspark_test.desc" % tmp_dir
...     with open(desc_file_path, "wb") as f:
...         _ = f.write(bytearray.fromhex(desc_hex))
...         f.flush()
...         message_name = 'SimpleMessage'
...         proto_df = df.select(to_protobuf(df.value,
...         desc_file_path, message_name).alias("value"))
...         proto_df.show(truncate=False)
...         proto_df = proto_df.select(from_protobuf(proto_df.value,
...         desc_file_path, message_name).alias("value"))
...         proto_df.show(truncate=False)
    +----------------------------------------+
    |value                                   |
    +----------------------------------------+
    |[08 02 12 05 41 6C 69 63 65 18 90 D5 06]|
    +----------------------------------------+
    +------------------+
    |value             |
    +------------------+
    |{2, Alice, 109200}|
    +------------------+

Tests Covered

• from_protobuf / to_protobuf (functions.py)

[SandishKumarHN]

@gengliangwang @rangadi here is the PR for pyspark functions for from_protobuf and to_protobuf

[zhengruifeng]

you may also add it to dev/sparktestsupport/modules.py

also cc @HyukjinKwon

[SandishKumarHN]

you may also add it to dev/sparktestsupport/modules.py

also cc @HyukjinKwon
@zhengruifeng added protobuf details to dev/sparktestsupport/modules.py

[HyukjinKwon]

Was this generated by something else? Should better have license header.

[SandishKumarHN]

@HyukjinKwon this is a binary file generated by the below commands(Protoc), though the .proto file has comments Protoc seems to strip down the comments. we can't manually edit the .desc file also. not sure what else options we have.

// protoc --java_out=connector/protobuf/src/test/resources/protobuf/ connector/protobuf/src/test/resources/protobuf/pyspark_test.proto
// protoc --descriptor_set_out=connector/protobuf/src/test/resources/protobuf/pyspark_test.desc --java_out=connector/protobuf/src/test/resources/protobuf/org/apache/spark/sql/protobuf/ connector/protobuf/src/test/resources/protobuf/pyspark_test.proto

[rangadi]

This could be part of build for tests, similar to how we generate compiled classes during build (same for Scala tests too). We can do that as a follow up with TODO.

[HyukjinKwon]

I wonder if we can emboded this into the doctests so users can copy and paste to test.

[SandishKumarHN]

@HyukjinKwon since from_protobuf and to_protobuf only take file path as a schema param, users have to pass pyspark_test.desc file. we have created multiple JIRAs on "Protobuf Support" EPIC to support multiple ways for passing protobuf schema.

[HyukjinKwon]

See https://github.com/apache/spark/pull/38212/files#r995302363

[HyukjinKwon]

Please also fix .github/labeler.yml to add this file into Protobuf compoenent.

[SandishKumarHN]

@HyukjinKwon @gengliangwang can we merge the PR?

[rangadi]

Is the desc file included in this PR? If not should this have failed?

[rangadi]

Nm. You are using the bytes.
Will generate it at build time in a follow up.

[HyukjinKwon]

LGTM. Thanks for working on this, this is a nice feature.

[HyukjinKwon]

Will merge once the tests pass.

[SandishKumarHN]

Will merge once the tests pass.

@HyukjinKwon All checks have passed.

[HyukjinKwon]

Merged to master.