[SPARK-55243][CONNECT] Allow setting binary headers via the -bin suffix in the Scala Connect client

### What changes were proposed in this pull request?
Automatically use the `Metadata.BINARY_BYTE_MARSHALLER` for `-bin` suffixed header keys, assuming base64-encoded header value strings set through the Scala Spark Connect client builder.

### Why are the changes needed?
The Scala Spark Connect client currently only allows setting `Metadata.ASCII_STRING_MARSHALLER` headers and fails if one tries to put a (binary) header with `-bin` key suffix:
```
[info]   java.lang.IllegalArgumentException: ASCII header is named test-bin.  Only binary headers may end with -bin
[info]   at com.google.common.base.Preconditions.checkArgument(Preconditions.java:445)
[info]   at io.grpc.Metadata$AsciiKey.<init>(Metadata.java:972)
[info]   at io.grpc.Metadata$AsciiKey.<init>(Metadata.java:966)
[info]   at io.grpc.Metadata$Key.of(Metadata.java:708)
[info]   at io.grpc.Metadata$Key.of(Metadata.java:704)
[info]   at org.apache.spark.sql.connect.client.SparkConnectClient$MetadataHeaderClientInterceptor$$anon$2.$anonfun$start$1(SparkConnectClient.scala:1112)
[info]   at org.apache.spark.sql.connect.client.SparkConnectClient$MetadataHeaderClientInterceptor$$anon$2.$anonfun$start$1$adapted(SparkConnectClient.scala:1106)
[info]   at scala.collection.immutable.Map$Map1.foreach(Map.scala:278)
[info]   at org.apache.spark.sql.connect.client.SparkConnectClient$MetadataHeaderClientInterceptor$$anon$2.start(SparkConnectClient.scala:1106)
[info]   at io.grpc.stub.ClientCalls.startCall(ClientCalls.java:435)
```

### Does this PR introduce _any_ user-facing change?
Current behaviour: Fails for all header key-value pairs if the key has the `-bin` suffix with an `IllegalArgumentException`.

New behaviour: Adds a `Metadata.BINARY_BYTE_MARSHALLER` header if the key has a `-bin` suffix and the value string is base64-encoded.

### How was this patch tested?
Added a test to `SparkConnectClientSuite`.

### Was this patch authored or co-authored using generative AI tooling?
No.

Closes apache#54016 from dillitz/fix-bin-header.

Authored-by: Robert Dillitz <r.dillitz@gmail.com>
Signed-off-by: Herman van Hövell <herman@databricks.com>

Author: dillitz

sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala

@@ -16,13 +16,14 @@
  */
 package org.apache.spark.sql.connect.client
 
-import java.util.UUID
+import java.nio.charset.StandardCharsets.UTF_8
+import java.util.{Base64, UUID}
 import java.util.concurrent.TimeUnit
 
 import scala.collection.mutable
 import scala.jdk.CollectionConverters._
 
-import io.grpc.{CallOptions, Channel, ClientCall, ClientInterceptor, MethodDescriptor, Server, Status, StatusRuntimeException}
+import io.grpc.{CallOptions, Channel, ClientCall, ClientInterceptor, Metadata, MethodDescriptor, Server, ServerCall, ServerCallHandler, ServerInterceptor, Status, StatusRuntimeException}
 import io.grpc.netty.NettyServerBuilder
 import io.grpc.stub.StreamObserver
 import org.scalatest.concurrent.Eventually
@@ -42,12 +43,13 @@ class SparkConnectClientSuite extends ConnectFunSuite {
   private var service: DummySparkConnectService = _
   private var server: Server = _
 
-  private def startDummyServer(port: Int): Unit = {
+  private def startDummyServer(port: Int, interceptors: Seq[ServerInterceptor] = Seq()): Unit = {
     service = new DummySparkConnectService
-    server = NettyServerBuilder
+    val serverBuilder = NettyServerBuilder
       .forPort(port)
       .addService(service)
-      .build()
+    interceptors.foreach(serverBuilder.intercept)
+    server = serverBuilder.build()
     server.start()
   }
 
@@ -622,6 +624,72 @@ class SparkConnectClientSuite extends ConnectFunSuite {
     // The client should try to fetch the config only once.
     assert(service.getAndClearLatestConfigRequests().size == 1)
   }
+
+  test("SPARK-55243: Binary headers use the correct marshaller") {
+    class HeadersInterceptor extends ServerInterceptor {
+      var headers: Option[Metadata] = None
+
+      override def interceptCall[ReqT, RespT](
+          call: ServerCall[ReqT, RespT],
+          headers: Metadata,
+          next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+        this.headers = Some(headers)
+        next.startCall(call, headers)
+      }
+    }
+
+    def buildClientWithHeader(key: String, value: String): SparkConnectClient = {
+      SparkConnectClient
+        .builder()
+        .connectionString(s"sc://localhost:${server.getPort}")
+        .option(key, value)
+        .build()
+    }
+
+    val headerInterceptor = new HeadersInterceptor()
+    startDummyServer(0, Seq(headerInterceptor))
+
+    val keyName = "test-bin"
+    val key = Metadata.Key.of(keyName, Metadata.BINARY_BYTE_MARSHALLER)
+    val binaryData = "test-binary-data"
+    val base64EncodedValue = Base64.getEncoder.encodeToString(binaryData.getBytes(UTF_8))
+
+    val plan = buildPlan("select * from range(10)")
+
+    // Successfully set and use base64-encoded -bin key.
+    client = buildClientWithHeader(keyName, base64EncodedValue)
+    client.execute(plan)
+
+    Eventually.eventually(timeout(5.seconds)) {
+      assert(headerInterceptor.headers.exists(_.containsKey(key)))
+      val bytes = headerInterceptor.headers.get.get(key)
+      assert(new String(bytes, UTF_8) == binaryData)
+    }
+
+    // Non base64-encoded -bin header throws IllegalArgumentException.
+    client = buildClientWithHeader(keyName, binaryData)
+
+    assertThrows[IllegalArgumentException] {
+      client.execute(plan)
+    }
+
+    // Non -bin headers keep using the ASCII marshaller.
+    val asciiKeyName = "test"
+    val asciiKey = Metadata.Key.of(asciiKeyName, Metadata.ASCII_STRING_MARSHALLER)
+
+    headerInterceptor.headers = None // Reset captured headers.
+
+    client = buildClientWithHeader(asciiKeyName, base64EncodedValue)
+    client.execute(plan)
+
+    Eventually.eventually(timeout(5.seconds)) {
+      assert(headerInterceptor.headers.exists(_.containsKey(asciiKey)))
+      val value = headerInterceptor.headers.get.get(asciiKey)
+      assert(value == base64EncodedValue)
+      // No BINARY_BYTE_MARSHALLER header.
+      assert(!headerInterceptor.headers.exists(_.containsKey(key)))
+    }
+  }
 }
 
 class DummySparkConnectService() extends SparkConnectServiceGrpc.SparkConnectServiceImplBase {

sql/connect/common/src/main/scala/org/apache/spark/sql/connect/client/SparkConnectClient.scala

@@ -18,7 +18,8 @@
 package org.apache.spark.sql.connect.client
 
 import java.net.URI
-import java.util.{Locale, UUID}
+import java.nio.charset.StandardCharsets.UTF_8
+import java.util.{Base64, Locale, UUID}
 import java.util.concurrent.Executor
 
 import scala.collection.mutable
@@ -1093,6 +1094,8 @@ object SparkConnectClient {
    */
   private[client] class MetadataHeaderClientInterceptor(metadata: Map[String, String])
       extends ClientInterceptor {
+    metadata.foreach { case (key, value) => assert(key != null && value != null) }
+
     override def interceptCall[ReqT, RespT](
         method: MethodDescriptor[ReqT, RespT],
         callOptions: CallOptions,
@@ -1103,7 +1106,13 @@ object SparkConnectClient {
             responseListener: ClientCall.Listener[RespT],
             headers: Metadata): Unit = {
           metadata.foreach { case (key, value) =>
-            headers.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value)
+            if (key.endsWith(Metadata.BINARY_HEADER_SUFFIX)) {
+              // Expects a base64-encoded value string.
+              val valueByteArray = Base64.getDecoder.decode(value.getBytes(UTF_8))
+              headers.put(Metadata.Key.of(key, Metadata.BINARY_BYTE_MARSHALLER), valueByteArray)
+            } else {
+              headers.put(Metadata.Key.of(key, Metadata.ASCII_STRING_MARSHALLER), value)
+            }
           }
           super.start(responseListener, headers)
         }