You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ya...@apache.org on 2023/05/18 02:06:22 UTC
[spark] branch master updated: [SPARK-43022][CONNECT] Support protobuf functions for Scala client
This is an automated email from the ASF dual-hosted git repository.
yangjie01 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 2dd112ec137 [SPARK-43022][CONNECT] Support protobuf functions for Scala client
2dd112ec137 is described below
commit 2dd112ec1372346ff573e5c273cc5cb93bda033a
Author: yangjie01 <ya...@baidu.com>
AuthorDate: Thu May 18 10:06:05 2023 +0800
[SPARK-43022][CONNECT] Support protobuf functions for Scala client
### What changes were proposed in this pull request?
This pr aims to support protobuf functions for Scala client.
### Why are the changes needed?
Add Spark connect jvm client api coverage.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
- Add new test
- Checked Scala 2.13
Closes #40654 from LuciferYang/protobuf-functions.
Lead-authored-by: yangjie01 <ya...@baidu.com>
Co-authored-by: YangJie <ya...@baidu.com>
Signed-off-by: yangjie01 <ya...@baidu.com>
---
.../org/apache/spark/sql/protobuf/functions.scala | 202 +++++++++++++++++++++
.../org/apache/spark/sql/FunctionTestSuite.scala | 11 ++
.../apache/spark/sql/PlanGenerationTestSuite.scala | 63 +++++++
.../src/test/resources/protobuf-tests/common.desc | 12 ++
.../from_protobuf_messageClassName.explain | 2 +
..._protobuf_messageClassName_descFilePath.explain | 2 +
...f_messageClassName_descFilePath_options.explain | 2 +
.../from_protobuf_messageClassName_options.explain | 2 +
.../to_protobuf_messageClassName.explain | 2 +
..._protobuf_messageClassName_descFilePath.explain | 2 +
...f_messageClassName_descFilePath_options.explain | 2 +
.../to_protobuf_messageClassName_options.explain | 2 +
.../queries/from_protobuf_messageClassName.json | 29 +++
.../from_protobuf_messageClassName.proto.bin | Bin 0 -> 125 bytes
...rom_protobuf_messageClassName_descFilePath.json | 33 ++++
...rotobuf_messageClassName_descFilePath.proto.bin | Bin 0 -> 156 bytes
...obuf_messageClassName_descFilePath_options.json | 46 +++++
...messageClassName_descFilePath_options.proto.bin | Bin 0 -> 206 bytes
.../from_protobuf_messageClassName_options.json | 42 +++++
...rom_protobuf_messageClassName_options.proto.bin | Bin 0 -> 174 bytes
.../queries/to_protobuf_messageClassName.json | 29 +++
.../queries/to_protobuf_messageClassName.proto.bin | Bin 0 -> 123 bytes
.../to_protobuf_messageClassName_descFilePath.json | 33 ++++
...rotobuf_messageClassName_descFilePath.proto.bin | Bin 0 -> 154 bytes
...obuf_messageClassName_descFilePath_options.json | 46 +++++
...messageClassName_descFilePath_options.proto.bin | Bin 0 -> 204 bytes
.../to_protobuf_messageClassName_options.json | 42 +++++
.../to_protobuf_messageClassName_options.proto.bin | Bin 0 -> 172 bytes
connector/connect/server/pom.xml | 6 +
.../sql/connect/planner/SparkConnectPlanner.scala | 59 ++++++
.../sql/protobuf/CatalystDataToProtobuf.scala | 2 +-
.../sql/protobuf/ProtobufDataToCatalyst.scala | 2 +-
32 files changed, 671 insertions(+), 2 deletions(-)
diff --git a/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/protobuf/functions.scala b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/protobuf/functions.scala
new file mode 100644
index 00000000000..c42f8417155
--- /dev/null
+++ b/connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/protobuf/functions.scala
@@ -0,0 +1,202 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.spark.sql.protobuf
+
+import scala.collection.JavaConverters._
+
+import org.apache.spark.annotation.Experimental
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.functions.{fnWithOptions, lit}
+
+// scalastyle:off: object.name
+object functions {
+ // scalastyle:on: object.name
+
+ /**
+ * Converts a binary column of Protobuf format into its corresponding catalyst value. The
+ * Protobuf definition is provided through Protobuf <i>descriptor file</i>.
+ *
+ * @param data
+ * the binary column.
+ * @param messageName
+ * the protobuf message name to look for in descriptor file.
+ * @param descFilePath
+ * the protobuf descriptor file.
+ * @param options
+ * @since 3.5.0
+ */
+ @Experimental
+ def from_protobuf(
+ data: Column,
+ messageName: String,
+ descFilePath: String,
+ options: java.util.Map[String, String]): Column = {
+ fnWithOptions(
+ "from_protobuf",
+ options.asScala.iterator,
+ data,
+ lit(messageName),
+ lit(descFilePath))
+ }
+
+ /**
+ * Converts a binary column of Protobuf format into its corresponding catalyst value. The
+ * Protobuf definition is provided through Protobuf <i>descriptor file</i>.
+ *
+ * @param data
+ * the binary column.
+ * @param messageName
+ * the protobuf MessageName to look for in descriptor file.
+ * @param descFilePath
+ * the protobuf descriptor file.
+ * @since 3.5.0
+ */
+ @Experimental
+ def from_protobuf(data: Column, messageName: String, descFilePath: String): Column = {
+ Column.fn("from_protobuf", data, lit(messageName), lit(descFilePath))
+ }
+
+ /**
+ * Converts a binary column of Protobuf format into its corresponding catalyst value.
+ * `messageClassName` points to Protobuf Java class. The jar containing Java class should be
+ * shaded. Specifically, `com.google.protobuf.*` should be shaded to
+ * `org.sparkproject.spark_protobuf.protobuf.*`.
+ * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from
+ * Protobuf files.
+ *
+ * @param data
+ * the binary column.
+ * @param messageClassName
+ * The full name for Protobuf Java class. E.g. <code>com.example.protos.ExampleEvent</code>.
+ * The jar with these classes needs to be shaded as described above.
+ * @since 3.5.0
+ */
+ @Experimental
+ def from_protobuf(data: Column, messageClassName: String): Column = {
+ Column.fn("from_protobuf", data, lit(messageClassName))
+ }
+
+ /**
+ * Converts a binary column of Protobuf format into its corresponding catalyst value.
+ * `messageClassName` points to Protobuf Java class. The jar containing Java class should be
+ * shaded. Specifically, `com.google.protobuf.*` should be shaded to
+ * `org.sparkproject.spark_protobuf.protobuf.*`.
+ * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from
+ * Protobuf files.
+ *
+ * @param data
+ * the binary column.
+ * @param messageClassName
+ * The full name for Protobuf Java class. E.g. <code>com.example.protos.ExampleEvent</code>.
+ * The jar with these classes needs to be shaded as described above.
+ * @param options
+ * @since 3.5.0
+ */
+ @Experimental
+ def from_protobuf(
+ data: Column,
+ messageClassName: String,
+ options: java.util.Map[String, String]): Column = {
+ fnWithOptions("from_protobuf", options.asScala.iterator, data, lit(messageClassName))
+ }
+
+ /**
+ * Converts a column into binary of protobuf format. The Protobuf definition is provided through
+ * Protobuf <i>descriptor file</i>.
+ *
+ * @param data
+ * the data column.
+ * @param messageName
+ * the protobuf MessageName to look for in descriptor file.
+ * @param descFilePath
+ * the protobuf descriptor file.
+ * @since 3.5.0
+ */
+ @Experimental
+ def to_protobuf(data: Column, messageName: String, descFilePath: String): Column = {
+ Column.fn("to_protobuf", data, lit(messageName), lit(descFilePath))
+ }
+
+ /**
+ * Converts a column into binary of protobuf format. The Protobuf definition is provided through
+ * Protobuf <i>descriptor file</i>.
+ *
+ * @param data
+ * the data column.
+ * @param messageName
+ * the protobuf MessageName to look for in descriptor file.
+ * @param descFilePath
+ * the protobuf descriptor file.
+ * @param options
+ * @since 3.5.0
+ */
+ @Experimental
+ def to_protobuf(
+ data: Column,
+ messageName: String,
+ descFilePath: String,
+ options: java.util.Map[String, String]): Column = {
+ fnWithOptions(
+ "to_protobuf",
+ options.asScala.iterator,
+ data,
+ lit(messageName),
+ lit(descFilePath))
+ }
+
+ /**
+ * Converts a column into binary of protobuf format. `messageClassName` points to Protobuf Java
+ * class. The jar containing Java class should be shaded. Specifically, `com.google.protobuf.*`
+ * should be shaded to `org.sparkproject.spark_protobuf.protobuf.*`.
+ * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from
+ * Protobuf files.
+ *
+ * @param data
+ * the data column.
+ * @param messageClassName
+ * The full name for Protobuf Java class. E.g. <code>com.example.protos.ExampleEvent</code>.
+ * The jar with these classes needs to be shaded as described above.
+ * @since 3.5.0
+ */
+ @Experimental
+ def to_protobuf(data: Column, messageClassName: String): Column = {
+ Column.fn("to_protobuf", data, lit(messageClassName))
+ }
+
+ /**
+ * Converts a column into binary of protobuf format. `messageClassName` points to Protobuf Java
+ * class. The jar containing Java class should be shaded. Specifically, `com.google.protobuf.*`
+ * should be shaded to `org.sparkproject.spark_protobuf.protobuf.*`.
+ * https://github.com/rangadi/shaded-protobuf-classes is useful to create shaded jar from
+ * Protobuf files.
+ *
+ * @param data
+ * the data column.
+ * @param messageClassName
+ * The full name for Protobuf Java class. E.g. <code>com.example.protos.ExampleEvent</code>.
+ * The jar with these classes needs to be shaded as described above.
+ * @param options
+ * @since 3.5.0
+ */
+ @Experimental
+ def to_protobuf(
+ data: Column,
+ messageClassName: String,
+ options: java.util.Map[String, String]): Column = {
+ fnWithOptions("to_protobuf", options.asScala.iterator, data, lit(messageClassName))
+ }
+}
diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/FunctionTestSuite.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/FunctionTestSuite.scala
index 9fdf0c32522..c92f77e8144 100644
--- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/FunctionTestSuite.scala
+++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/FunctionTestSuite.scala
@@ -18,9 +18,12 @@ package org.apache.spark.sql
import java.util.Collections
+import scala.collection.JavaConverters._
+
import org.apache.spark.sql.avro.{functions => avroFn}
import org.apache.spark.sql.connect.client.util.ConnectFunSuite
import org.apache.spark.sql.functions._
+import org.apache.spark.sql.protobuf.{functions => pbFn}
import org.apache.spark.sql.types.{DataType, StructType}
/**
@@ -233,6 +236,14 @@ class FunctionTestSuite extends ConnectFunSuite {
a,
"""{"type": "int", "name": "id"}""",
Collections.emptyMap[String, String]))
+ testEquals(
+ "from_protobuf",
+ pbFn.from_protobuf(a, "FakeMessage", "fakePath", Map.empty[String, String].asJava),
+ pbFn.from_protobuf(a, "FakeMessage", "fakePath"))
+ testEquals(
+ "to_protobuf",
+ pbFn.to_protobuf(a, "FakeMessage", "fakePath", Map.empty[String, String].asJava),
+ pbFn.to_protobuf(a, "FakeMessage", "fakePath"))
test("assert_true no message") {
val e = assert_true(a).expr
diff --git a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
index b9b776420c2..7ece54d0439 100644
--- a/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
+++ b/connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
@@ -30,6 +30,7 @@ import io.grpc.inprocess.InProcessChannelBuilder
import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
import org.apache.spark.connect.proto
+import org.apache.spark.connect.proto.StorageLevel
import org.apache.spark.internal.Logging
import org.apache.spark.sql.{functions => fn}
import org.apache.spark.sql.avro.{functions => avroFn}
@@ -39,6 +40,7 @@ import org.apache.spark.sql.connect.client.SparkConnectClient
import org.apache.spark.sql.connect.client.util.ConnectFunSuite
import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.lit
+import org.apache.spark.sql.protobuf.{functions => pbFn}
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.CalendarInterval
@@ -2242,4 +2244,65 @@ class PlanGenerationTestSuite
test("to_avro without schema") {
simple.select(avroFn.to_avro(fn.col("id")))
}
+
+ /* Protobuf functions */
+ // scalastyle:off line.size.limit
+ // If `common.desc` needs to be updated, execute the following command to regenerate it:
+ // 1. cd connector/connect/common/src/main/protobuf/spark/connect
+ // 2. protoc --include_imports --descriptor_set_out=../../../../test/resources/protobuf-tests/common.desc common.proto
+ // scalastyle:on line.size.limit
+ private val testDescFilePath: String = java.nio.file.Paths
+ .get("../", "common", "src", "test", "resources", "protobuf-tests")
+ .resolve("common.desc")
+ .toString
+
+ test("from_protobuf messageClassName") {
+ binary.select(pbFn.from_protobuf(fn.col("bytes"), classOf[StorageLevel].getName))
+ }
+
+ test("from_protobuf messageClassName options") {
+ binary.select(
+ pbFn.from_protobuf(
+ fn.col("bytes"),
+ classOf[StorageLevel].getName,
+ Map("recursive.fields.max.depth" -> "2").asJava))
+ }
+
+ test("from_protobuf messageClassName descFilePath options") {
+ binary.select(
+ pbFn.from_protobuf(
+ fn.col("bytes"),
+ "StorageLevel",
+ testDescFilePath,
+ Map("recursive.fields.max.depth" -> "2").asJava))
+ }
+
+ test("from_protobuf messageClassName descFilePath") {
+ binary.select(pbFn.from_protobuf(fn.col("bytes"), "StorageLevel", testDescFilePath))
+ }
+
+ test("to_protobuf messageClassName") {
+ binary.select(pbFn.to_protobuf(fn.col("bytes"), classOf[StorageLevel].getName))
+ }
+
+ test("to_protobuf messageClassName options") {
+ binary.select(
+ pbFn.to_protobuf(
+ fn.col("bytes"),
+ classOf[StorageLevel].getName,
+ Map("recursive.fields.max.depth" -> "2").asJava))
+ }
+
+ test("to_protobuf messageClassName descFilePath options") {
+ binary.select(
+ pbFn.to_protobuf(
+ fn.col("bytes"),
+ "StorageLevel",
+ testDescFilePath,
+ Map("recursive.fields.max.depth" -> "2").asJava))
+ }
+
+ test("to_protobuf messageClassName descFilePath") {
+ binary.select(pbFn.to_protobuf(fn.col("bytes"), "StorageLevel", testDescFilePath))
+ }
}
diff --git a/connector/connect/common/src/test/resources/protobuf-tests/common.desc b/connector/connect/common/src/test/resources/protobuf-tests/common.desc
new file mode 100644
index 00000000000..bec358e02bd
--- /dev/null
+++ b/connector/connect/common/src/test/resources/protobuf-tests/common.desc
@@ -0,0 +1,12 @@
+
+�
+ common.proto
spark.connect"�
+ StorageLevel
+use_disk (RuseDisk
+
+use_memory (R useMemory
+ use_off_heap (R
+useOffHeap"
+ deserialized (R deserialized
+ replication (R replicationB"
+ org.apache.spark.connect.protoPbproto3
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName.explain b/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName.explain
new file mode 100644
index 00000000000..e7a1867fe90
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName.explain
@@ -0,0 +1,2 @@
+Project [from_protobuf(bytes#0, org.apache.spark.connect.proto.StorageLevel, None) AS from_protobuf(bytes)#0]
++- LocalRelation <empty>, [id#0L, bytes#0]
diff --git a/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName_descFilePath.explain b/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName_descFilePath.explain
new file mode 100644
index 00000000000..b3eaa5f44d6
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName_descFilePath.explain
@@ -0,0 +1,2 @@
+Project [from_protobuf(bytes#0, StorageLevel, Some(../common/src/test/resources/protobuf-tests/common.desc)) AS from_protobuf(bytes)#0]
++- LocalRelation <empty>, [id#0L, bytes#0]
diff --git a/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName_descFilePath_options.explain b/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName_descFilePath_options.explain
new file mode 100644
index 00000000000..85851e14da3
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName_descFilePath_options.explain
@@ -0,0 +1,2 @@
+Project [from_protobuf(bytes#0, StorageLevel, Some(../common/src/test/resources/protobuf-tests/common.desc), (recursive.fields.max.depth,2)) AS from_protobuf(bytes)#0]
++- LocalRelation <empty>, [id#0L, bytes#0]
diff --git a/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName_options.explain b/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName_options.explain
new file mode 100644
index 00000000000..c02d829fcac
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/explain-results/from_protobuf_messageClassName_options.explain
@@ -0,0 +1,2 @@
+Project [from_protobuf(bytes#0, org.apache.spark.connect.proto.StorageLevel, None, (recursive.fields.max.depth,2)) AS from_protobuf(bytes)#0]
++- LocalRelation <empty>, [id#0L, bytes#0]
diff --git a/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName.explain b/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName.explain
new file mode 100644
index 00000000000..e7f70fa2c1a
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName.explain
@@ -0,0 +1,2 @@
+Project [to_protobuf(bytes#0, org.apache.spark.connect.proto.StorageLevel, None) AS to_protobuf(bytes)#0]
++- LocalRelation <empty>, [id#0L, bytes#0]
diff --git a/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName_descFilePath.explain b/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName_descFilePath.explain
new file mode 100644
index 00000000000..55b53cfcdb1
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName_descFilePath.explain
@@ -0,0 +1,2 @@
+Project [to_protobuf(bytes#0, StorageLevel, Some(../common/src/test/resources/protobuf-tests/common.desc)) AS to_protobuf(bytes)#0]
++- LocalRelation <empty>, [id#0L, bytes#0]
diff --git a/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName_descFilePath_options.explain b/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName_descFilePath_options.explain
new file mode 100644
index 00000000000..eafd4124d8f
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName_descFilePath_options.explain
@@ -0,0 +1,2 @@
+Project [to_protobuf(bytes#0, StorageLevel, Some(../common/src/test/resources/protobuf-tests/common.desc), (recursive.fields.max.depth,2)) AS to_protobuf(bytes)#0]
++- LocalRelation <empty>, [id#0L, bytes#0]
diff --git a/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName_options.explain b/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName_options.explain
new file mode 100644
index 00000000000..a5d8851a7d1
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/explain-results/to_protobuf_messageClassName_options.explain
@@ -0,0 +1,2 @@
+Project [to_protobuf(bytes#0, org.apache.spark.connect.proto.StorageLevel, None, (recursive.fields.max.depth,2)) AS to_protobuf(bytes)#0]
++- LocalRelation <empty>, [id#0L, bytes#0]
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName.json b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName.json
new file mode 100644
index 00000000000..dc23ac2a117
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName.json
@@ -0,0 +1,29 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema": "struct\u003cid:bigint,bytes:binary\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "from_protobuf",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "bytes"
+ }
+ }, {
+ "literal": {
+ "string": "org.apache.spark.connect.proto.StorageLevel"
+ }
+ }]
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName.proto.bin b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName.proto.bin
new file mode 100644
index 00000000000..cc46234b747
Binary files /dev/null and b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName.proto.bin differ
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath.json b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath.json
new file mode 100644
index 00000000000..05e210434cd
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath.json
@@ -0,0 +1,33 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema": "struct\u003cid:bigint,bytes:binary\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "from_protobuf",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "bytes"
+ }
+ }, {
+ "literal": {
+ "string": "StorageLevel"
+ }
+ }, {
+ "literal": {
+ "string": "../common/src/test/resources/protobuf-tests/common.desc"
+ }
+ }]
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath.proto.bin b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath.proto.bin
new file mode 100644
index 00000000000..c4c2b2126b8
Binary files /dev/null and b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath.proto.bin differ
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath_options.json b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath_options.json
new file mode 100644
index 00000000000..ff747ede6f1
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath_options.json
@@ -0,0 +1,46 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema": "struct\u003cid:bigint,bytes:binary\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "from_protobuf",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "bytes"
+ }
+ }, {
+ "literal": {
+ "string": "StorageLevel"
+ }
+ }, {
+ "literal": {
+ "string": "../common/src/test/resources/protobuf-tests/common.desc"
+ }
+ }, {
+ "unresolvedFunction": {
+ "functionName": "map",
+ "arguments": [{
+ "literal": {
+ "string": "recursive.fields.max.depth"
+ }
+ }, {
+ "literal": {
+ "string": "2"
+ }
+ }]
+ }
+ }]
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath_options.proto.bin b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath_options.proto.bin
new file mode 100644
index 00000000000..640df1c8d1b
Binary files /dev/null and b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_descFilePath_options.proto.bin differ
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_options.json b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_options.json
new file mode 100644
index 00000000000..36f69646ef8
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_options.json
@@ -0,0 +1,42 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema": "struct\u003cid:bigint,bytes:binary\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "from_protobuf",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "bytes"
+ }
+ }, {
+ "literal": {
+ "string": "org.apache.spark.connect.proto.StorageLevel"
+ }
+ }, {
+ "unresolvedFunction": {
+ "functionName": "map",
+ "arguments": [{
+ "literal": {
+ "string": "recursive.fields.max.depth"
+ }
+ }, {
+ "literal": {
+ "string": "2"
+ }
+ }]
+ }
+ }]
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_options.proto.bin b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_options.proto.bin
new file mode 100644
index 00000000000..72a1c6b8207
Binary files /dev/null and b/connector/connect/common/src/test/resources/query-tests/queries/from_protobuf_messageClassName_options.proto.bin differ
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName.json b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName.json
new file mode 100644
index 00000000000..921c1b800a0
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName.json
@@ -0,0 +1,29 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema": "struct\u003cid:bigint,bytes:binary\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "to_protobuf",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "bytes"
+ }
+ }, {
+ "literal": {
+ "string": "org.apache.spark.connect.proto.StorageLevel"
+ }
+ }]
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName.proto.bin b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName.proto.bin
new file mode 100644
index 00000000000..5cc7c49882c
Binary files /dev/null and b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName.proto.bin differ
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath.json b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath.json
new file mode 100644
index 00000000000..e5bbb2ba57d
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath.json
@@ -0,0 +1,33 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema": "struct\u003cid:bigint,bytes:binary\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "to_protobuf",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "bytes"
+ }
+ }, {
+ "literal": {
+ "string": "StorageLevel"
+ }
+ }, {
+ "literal": {
+ "string": "../common/src/test/resources/protobuf-tests/common.desc"
+ }
+ }]
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath.proto.bin b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath.proto.bin
new file mode 100644
index 00000000000..346fc3d90f3
Binary files /dev/null and b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath.proto.bin differ
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath_options.json b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath_options.json
new file mode 100644
index 00000000000..183c8a25ce1
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath_options.json
@@ -0,0 +1,46 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema": "struct\u003cid:bigint,bytes:binary\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "to_protobuf",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "bytes"
+ }
+ }, {
+ "literal": {
+ "string": "StorageLevel"
+ }
+ }, {
+ "literal": {
+ "string": "../common/src/test/resources/protobuf-tests/common.desc"
+ }
+ }, {
+ "unresolvedFunction": {
+ "functionName": "map",
+ "arguments": [{
+ "literal": {
+ "string": "recursive.fields.max.depth"
+ }
+ }, {
+ "literal": {
+ "string": "2"
+ }
+ }]
+ }
+ }]
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath_options.proto.bin b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath_options.proto.bin
new file mode 100644
index 00000000000..cc4d08ee581
Binary files /dev/null and b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_descFilePath_options.proto.bin differ
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_options.json b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_options.json
new file mode 100644
index 00000000000..8787f0fc15d
--- /dev/null
+++ b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_options.json
@@ -0,0 +1,42 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema": "struct\u003cid:bigint,bytes:binary\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "to_protobuf",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "bytes"
+ }
+ }, {
+ "literal": {
+ "string": "org.apache.spark.connect.proto.StorageLevel"
+ }
+ }, {
+ "unresolvedFunction": {
+ "functionName": "map",
+ "arguments": [{
+ "literal": {
+ "string": "recursive.fields.max.depth"
+ }
+ }, {
+ "literal": {
+ "string": "2"
+ }
+ }]
+ }
+ }]
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_options.proto.bin b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_options.proto.bin
new file mode 100644
index 00000000000..9ef8348446a
Binary files /dev/null and b/connector/connect/common/src/test/resources/query-tests/queries/to_protobuf_messageClassName_options.proto.bin differ
diff --git a/connector/connect/server/pom.xml b/connector/connect/server/pom.xml
index 8313f21f4ba..f9dc40ce85a 100644
--- a/connector/connect/server/pom.xml
+++ b/connector/connect/server/pom.xml
@@ -111,6 +111,12 @@
<version>${project.version}</version>
<scope>provided</scope>
</dependency>
+ <dependency>
+ <groupId>org.apache.spark</groupId>
+ <artifactId>spark-protobuf_${scala.binary.version}</artifactId>
+ <version>${project.version}</version>
+ <scope>provided</scope>
+ </dependency>
<dependency>
<groupId>org.apache.spark</groupId>
<artifactId>spark-catalyst_${scala.binary.version}</artifactId>
diff --git a/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
index 307cf1d054d..eac2b339616 100644
--- a/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
+++ b/connector/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -69,6 +69,7 @@ import org.apache.spark.sql.execution.stat.StatFunctions
import org.apache.spark.sql.execution.streaming.StreamingQueryWrapper
import org.apache.spark.sql.expressions.ReduceAggregator
import org.apache.spark.sql.internal.{CatalogImpl, TypedAggUtils}
+import org.apache.spark.sql.protobuf.{CatalystDataToProtobuf, ProtobufDataToCatalyst}
import org.apache.spark.sql.streaming.{StreamingQuery, StreamingQueryProgress, Trigger}
import org.apache.spark.sql.types._
import org.apache.spark.sql.util.CaseInsensitiveStringMap
@@ -1448,6 +1449,51 @@ class SparkConnectPlanner(val session: SparkSession) {
*/
private def transformUnregisteredFunction(
fun: proto.Expression.UnresolvedFunction): Option[Expression] = {
+ def extractArgsOfProtobufFunction(
+ functionName: String,
+ argumentsCount: Int,
+ children: Seq[Expression]): (String, Option[String], Map[String, String]) = {
+ val messageClassName = children(1) match {
+ case Literal(s, StringType) if s != null => s.toString
+ case other =>
+ throw InvalidPlanInput(
+ s"MessageClassName in $functionName should be a literal string, but got $other")
+ }
+ val (descFilePathOpt, options) = if (argumentsCount == 2) {
+ (None, Map.empty[String, String])
+ } else if (argumentsCount == 3) {
+ children(2) match {
+ case Literal(s, StringType) if s != null =>
+ (Some(s.toString), Map.empty[String, String])
+ case UnresolvedFunction(Seq("map"), arguments, _, _, _) =>
+ (None, ExprUtils.convertToMapData(CreateMap(arguments)))
+ case other =>
+ throw InvalidPlanInput(
+ s"The valid type for the 3rd arg in $functionName is string or map, but got $other")
+ }
+ } else if (argumentsCount == 4) {
+ val filePathOpt = children(2) match {
+ case Literal(s, StringType) if s != null =>
+ Some(s.toString)
+ case other =>
+ throw InvalidPlanInput(
+ s"DescFilePath in $functionName should be a literal string, but got $other")
+ }
+ val map = children(3) match {
+ case UnresolvedFunction(Seq("map"), arguments, _, _, _) =>
+ ExprUtils.convertToMapData(CreateMap(arguments))
+ case other =>
+ throw InvalidPlanInput(
+ s"Options in $functionName should be created by map, but got $other")
+ }
+ (filePathOpt, map)
+ } else {
+ throw InvalidPlanInput(
+ s"$functionName requires 2 ~ 4 arguments, but got $argumentsCount ones!")
+ }
+ (messageClassName, descFilePathOpt, options)
+ }
+
fun.getFunctionName match {
case "product" if fun.getArgumentsCount == 1 =>
Some(
@@ -1600,6 +1646,19 @@ class SparkConnectPlanner(val session: SparkSession) {
val expr = transformExpression(fun.getArguments(0))
Some(transformUnregisteredUDF(MLFunctions.arrayToVectorUdf, Seq(expr)))
+ // Protobuf-specific functions
+ case "from_protobuf" if Seq(2, 3, 4).contains(fun.getArgumentsCount) =>
+ val children = fun.getArgumentsList.asScala.toSeq.map(transformExpression)
+ val (messageClassName, descFilePathOpt, options) =
+ extractArgsOfProtobufFunction("from_protobuf", fun.getArgumentsCount, children)
+ Some(ProtobufDataToCatalyst(children.head, messageClassName, descFilePathOpt, options))
+
+ case "to_protobuf" if Seq(2, 3, 4).contains(fun.getArgumentsCount) =>
+ val children = fun.getArgumentsList.asScala.toSeq.map(transformExpression)
+ val (messageClassName, descFilePathOpt, options) =
+ extractArgsOfProtobufFunction("to_protobuf", fun.getArgumentsCount, children)
+ Some(CatalystDataToProtobuf(children.head, messageClassName, descFilePathOpt, options))
+
case _ => None
}
}
diff --git a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/CatalystDataToProtobuf.scala b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/CatalystDataToProtobuf.scala
index 12561fe51e6..62d251fb869 100644
--- a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/CatalystDataToProtobuf.scala
+++ b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/CatalystDataToProtobuf.scala
@@ -23,7 +23,7 @@ import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCo
import org.apache.spark.sql.protobuf.utils.ProtobufUtils
import org.apache.spark.sql.types.{BinaryType, DataType}
-private[protobuf] case class CatalystDataToProtobuf(
+private[sql] case class CatalystDataToProtobuf(
child: Expression,
messageName: String,
descFilePath: Option[String] = None,
diff --git a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDataToCatalyst.scala b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDataToCatalyst.scala
index 184e9380124..128169f46a2 100644
--- a/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDataToCatalyst.scala
+++ b/connector/protobuf/src/main/scala/org/apache/spark/sql/protobuf/ProtobufDataToCatalyst.scala
@@ -29,7 +29,7 @@ import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors
import org.apache.spark.sql.protobuf.utils.{ProtobufOptions, ProtobufUtils, SchemaConverters}
import org.apache.spark.sql.types.{AbstractDataType, BinaryType, DataType, StructType}
-private[protobuf] case class ProtobufDataToCatalyst(
+private[sql] case class ProtobufDataToCatalyst(
child: Expression,
messageName: String,
descFilePath: Option[String] = None,
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org