You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "xi-db (via GitHub)" <gi...@apache.org> on 2026/02/03 09:56:41 UTC
[PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
xi-db opened a new pull request, #54113:
URL: https://github.com/apache/spark/pull/54113
### What changes were proposed in this pull request?
Previously, decompression logic was scattered across multiple handlers (AnalyzePlanHandler, ExecutePlanHandler, and probably other handlers in the future), with each handler responsible for decompressing its own compressed requests. This approach made it difficult to maintain consistent decompression behavior across different RPC types.
This PR introduces RequestDecompressionInterceptor as a gRPC interceptor that handles decompression for all compressed requests before they reach downstream handlers.
### Why are the changes needed?
This refactoring ensures single responsibility and consistency, as decompression logic is now in one place, making it easier to maintain.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
New tests for the interceptor and existing tests for plan compression.
### Was this patch authored or co-authored using generative AI tooling?
No.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792287908
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Yes, the pattern matching of request.getAnalyzeCase cannot be avoided. But my concern is that decompress is called in every case, sometimes multiple times, which seems a bit redundant and makes the code somewhat hard to read.
Below is a rough sketch (generated by ChatGPT).
```scala
private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
: (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
val userId = request.getUserContext.getUserId
val sessionId = request.getSessionId
val planAccessors: Seq[(proto.AnalyzePlanRequest => proto.Plan,
(proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest)] =
request.getAnalyzeCase match {
case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
Seq(
(_.getSchema.getPlan, (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
Seq(
(_.getExplain.getPlan, (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
Seq(
(_.getSameSemantics.getTargetPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build()),
(_.getSameSemantics.getOtherPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
)
// other single-plan cases...
case _ => Seq()
}
// Decompress all plans in sequence
var req = request
var sizes: Seq[Option[Long]] = Seq.empty
planAccessors.foreach { case (getter, setter) =>
val (newReq, size) = decompressPlanGeneric(req, userId, sessionId, getter, setter)
req = newReq
sizes = sizes :+ size
}
val targetSize = sizes.headOption.flatten
val otherSize = if (sizes.length > 1) sizes(1) else None
(req, targetSize, otherSize)
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2794946393
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Hmm, another benefit I can foresee is that if other developers need to add more getters/setters for a new AnalyzeCase, it will add less cognitive load to update the code if we do the refactoring. Otherwise, they would still need to call decompress for reqWithTarget and finalReq (for the second plan), and pass the size/other size (for the second plan), instead of just figuring out how to write the getter/setter.
@hvanhovell @vicennial WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2770299648
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
The name seems ad hoc and this naming pattern may be difficult to extend for a third plan.(other-other-compressed-size?). Would it be possible to store an array of compressed sizes?
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
The name seems ad hoc and this naming pattern may be difficult to extend for a third plan(other-other-compressed-size?). Would it be possible to store an array of compressed sizes?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2794654802
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
I agree this refactoring will be better if many cases have many plans, but in fact, all the cases except the SameSemantics have only one plan. I think this approach will make it a bit harder to understand.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792287908
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Yes, the pattern matching in request.getAnalyzeCase cannot be avoided. However, my concern is that decompress is called in many cases—sometimes multiple times—which seems a bit redundant and makes the code somewhat hard to read.
Below is a rough sketch (generated by ChatGPT) to help reduce duplication:
```scala
private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
: (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
val userId = request.getUserContext.getUserId
val sessionId = request.getSessionId
val planAccessors: Seq[(proto.AnalyzePlanRequest => proto.Plan,
(proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest)] =
request.getAnalyzeCase match {
case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
Seq(
(_.getSchema.getPlan, (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
Seq(
(_.getExplain.getPlan, (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
Seq(
(_.getSameSemantics.getTargetPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build()),
(_.getSameSemantics.getOtherPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
)
// other single-plan cases...
}
// Decompress all plans in sequence
var req = request
var sizes: Seq[Option[Long]] = Seq.empty
planAccessors.foreach { case (getter, setter) =>
val (newReq, size) = decompressPlanGeneric(req, userId, sessionId, getter, setter)
req = newReq
sizes = sizes :+ size
}
val targetSize = sizes.headOption.flatten
val otherSize = if (sizes.length > 1) sizes(1) else None
(req, targetSize, otherSize)
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2794946393
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Hmm, another benefit I can foresee is that if other developers need to add getters/setters for a new AnalyzeCase, it will create less cognitive load when updating the code if we do the refactoring. Otherwise, they would still need to call decompress for reqWithTarget and finalReq (for the second plan if needed), and pass the size/other size (for the second plan if needed), instead of just figuring out how to write the getter/setter.
@hvanhovell @vicennial WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2794946393
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Hmm, another benefit I can foresee is that if other developers need to add getters/setters for a new AnalyzeCase, it will create less cognitive load when updating the code if we do the refactoring. Otherwise, they would still need to call decompress for reqWithTarget and finalReq (for the second plan, if needed), and pass the size/other size (for the second plan, if needed), instead of just figuring out how to implement the getter/setter.
@hvanhovell @vicennial WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2878466512
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAnalyzeHandler.scala:
##########
@@ -69,9 +70,23 @@ private[connect] class SparkConnectAnalyzeHandler(
throw new AnalysisException("ATTEMPT_ANALYSIS_IN_PIPELINE_QUERY_FUNCTION", Map())
}
+ // Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
+ val compressedSize = RequestDecompressionContext.getCompressedSize
+ val otherCompressedSize = RequestDecompressionContext.getOtherCompressedSize
+ if (compressedSize.isDefined || otherCompressedSize.isDefined) {
+ logDebug(
+ log"AnalyzePlan request received with compressed plan: " +
+ log"compressedSize=${MDC(BYTE_SIZE, compressedSize.getOrElse(0L))} bytes" +
+ otherCompressedSize
+ .map { size =>
+ log", otherCompressedSize=${MDC(BYTE_SIZE, size)} bytes"
+ }
+ .getOrElse(log""))
+ }
+
def transformRelation(rel: proto.Relation) = planner.transformRelation(rel, cachePlan = true)
def transformRelationPlan(plan: proto.Plan) = {
Review Comment:
Yes, but I would like to keep it because if we want to have plan-level operations in the future we will need this tiny helper back. Changed this to one-line `def transformRelationPlan(plan: proto.Plan) = transformRelation(plan.getRoot)` to be simply.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2770299648
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
The name seems ad hoc and may be difficult to extend for a third plan (other-other-compressed-size?). Would it be possible to store an array of compressed sizes?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792393799
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
Not very important, but I still don’t follow this naming pattern (the first one should be named target-compressed-size if we follow it strictly)... That said, I still feel that storing the sizes in an array, if possible, could simplify the logging logic below:
```scala
// Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
val compressedSize = RequestDecompressionContext.getCompressedSize
val otherCompressedSize = RequestDecompressionContext.getOtherCompressedSize
if (compressedSize.isDefined || otherCompressedSize.isDefined) {
logDebug(
log"AnalyzePlan request received with compressed plan: " +
log"compressedSize=${MDC(BYTE_SIZE, compressedSize.getOrElse(0L))} bytes" +
otherCompressedSize
.map { size =>
log", otherCompressedSize=${MDC(BYTE_SIZE, size)} bytes"
}
.getOrElse(log""))
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2794946393
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Another benefit I can foresee is that if other developers need to add more getters/setters for a new AnalyzeCase, it will add less cognitive load to update the code if we do the refactoring. Otherwise, they would still need to call decompress for reqWithTarget and finalReq (for the second plan), and pass the size/other size (for the second plan), instead of just figuring out how to write the getter/setter.
@hvanhovell @vicennial WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2787688891
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
Actually the name comes from the Spark Connect protocol: the two values correspond to target_plan and other_plan: https://github.com/apache/spark/blob/ebd5b007fcf203eadcf8b037ab2b99577490f869/sql/connect/common/src/main/protobuf/spark/connect/base.proto#L190-L196
I don't think it will be extented to 3 plans, because almost all the APIs in Spark Connect protocol has only one plan.
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
Actually the name comes from the Spark Connect protocol: the two values correspond to target_plan and other_plan: https://github.com/apache/spark/blob/ebd5b007fcf203eadcf8b037ab2b99577490f869/sql/connect/common/src/main/protobuf/spark/connect/base.proto#L190-L196
I don't think it will be extented to 3 plans, because almost all the APIs in Spark Connect protocol have only one plan.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #54113:
URL: https://github.com/apache/spark/pull/54113#issuecomment-3840291561
## JIRA Issue Information
=== Improvement SPARK-55338 ===
Summary: Centralize request decompression logic in gRPC interceptor
Assignee: None
Status: Open
Affected: ["4.1.1"]
---
*This comment was automatically generated by GitHub Actions*
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2787716706
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Yes, the existing code actually follows exactly the same idea, where in each case, we call the internal helper `decompress`, which takes the planGetter and planSetter, and the helper method already extracts the duplicate logic. But the pattern matching of request.getAnalyzeCase cannot be avoided :(
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2794946393
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Hmm, another benefit I can foresee is that if other developers need to add more getters/setters for a new AnalyzeCase, it will create less cognitive load when updating the code if we do the refactoring. Otherwise, they would still need to call decompress for reqWithTarget and finalReq (for the second plan), and pass the size/other size (for the second plan), instead of just figuring out how to write the getter/setter.
@hvanhovell @vicennial WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2864997071
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectExecutePlanHandler.scala:
##########
@@ -36,20 +36,20 @@ class SparkConnectExecutePlanHandler(responseObserver: StreamObserver[proto.Exec
.getOrCreateIsolatedSession(v.getUserContext.getUserId, v.getSessionId, previousSessionId)
val executeKey = ExecuteKey(v, sessionHolder)
- val decompressedRequest =
- v.toBuilder.setPlan(PlanCompressionUtils.decompressPlan(v.getPlan)).build()
+ // Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
+ val compressedSize = RequestDecompressionContext.getCompressedSize
Review Comment:
If we are only going to log this, why have a RequestDecompressionContext at all? You also log in RequestDecompressionInterceptor.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "asf-gitbox-commits (via GitHub)" <gi...@apache.org>.
asf-gitbox-commits closed pull request #54113: [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor
URL: https://github.com/apache/spark/pull/54113
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792393799
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
I am still a bit confused about the naming pattern (the first one should be named target-compressed-size if we follow it strictly). And I still feel that storing the sizes in an array (if possible) could simplify the below logging logic:
```scala
// Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
val compressedSize = RequestDecompressionContext.getCompressedSize
val otherCompressedSize = RequestDecompressionContext.getOtherCompressedSize
if (compressedSize.isDefined || otherCompressedSize.isDefined) {
logDebug(
log"AnalyzePlan request received with compressed plan: " +
log"compressedSize=${MDC(BYTE_SIZE, compressedSize.getOrElse(0L))} bytes" +
otherCompressedSize
.map { size =>
log", otherCompressedSize=${MDC(BYTE_SIZE, size)} bytes"
}
.getOrElse(log""))
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792287908
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Yes, the pattern matching of request.getAnalyzeCase cannot be avoided. But my concern is that decompress is called in every case, sometimes multiple times, which seems a bit redundant and makes the code somewhat hard to read.
Below is a rough sketch (generated by ChatGPT).
```scala
private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
: (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
val userId = request.getUserContext.getUserId
val sessionId = request.getSessionId
val planAccessors: Seq[(proto.AnalyzePlanRequest => proto.Plan,
(proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest)] =
request.getAnalyzeCase match {
case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
Seq(
(_.getSchema.getPlan, (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
Seq(
(_.getExplain.getPlan, (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
Seq(
(_.getSameSemantics.getTargetPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build()),
(_.getSameSemantics.getOtherPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
)
// other single-plan cases...
}
// Decompress all plans in sequence
var req = request
var sizes: Seq[Option[Long]] = Seq.empty
planAccessors.foreach { case (getter, setter) =>
val (newReq, size) = decompressPlanGeneric(req, userId, sessionId, getter, setter)
req = newReq
sizes = sizes :+ size
}
val targetSize = sizes.headOption.flatten
val otherSize = if (sizes.length > 1) sizes(1) else None
(req, targetSize, otherSize)
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792393799
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
Not sure if I fully follow this naming pattern (the first one should be named target-compressed-size if we follow it strictly), but I still feel that storing the sizes in an array, if possible, could simplify the logging logic below:
```scala
// Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
val compressedSize = RequestDecompressionContext.getCompressedSize
val otherCompressedSize = RequestDecompressionContext.getOtherCompressedSize
if (compressedSize.isDefined || otherCompressedSize.isDefined) {
logDebug(
log"AnalyzePlan request received with compressed plan: " +
log"compressedSize=${MDC(BYTE_SIZE, compressedSize.getOrElse(0L))} bytes" +
otherCompressedSize
.map { size =>
log", otherCompressedSize=${MDC(BYTE_SIZE, size)} bytes"
}
.getOrElse(log""))
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "hvanhovell (via GitHub)" <gi...@apache.org>.
hvanhovell commented on PR #54113:
URL: https://github.com/apache/spark/pull/54113#issuecomment-3999603379
LGTM - merging to master.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2769607332
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
I was wondering if it makes sense to return Seq[(planGetter, planSetter)] and run the decompression logic elsewhere to reduce code duplication.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2770299648
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
The name seems ad hoc and this naming pattern may be difficult to extend for a third plan(other-other-compressed-size?). Would it be possible to store an array of compressed sizes instead?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792393799
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
Not very important, but I still don’t follow this naming pattern (the first one should be named target-compressed-size if we follow it strictly). That said, I still feel that storing the sizes in an array, if possible, could simplify the logging logic below:
```scala
// Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
val compressedSize = RequestDecompressionContext.getCompressedSize
val otherCompressedSize = RequestDecompressionContext.getOtherCompressedSize
if (compressedSize.isDefined || otherCompressedSize.isDefined) {
logDebug(
log"AnalyzePlan request received with compressed plan: " +
log"compressedSize=${MDC(BYTE_SIZE, compressedSize.getOrElse(0L))} bytes" +
otherCompressedSize
.map { size =>
log", otherCompressedSize=${MDC(BYTE_SIZE, size)} bytes"
}
.getOrElse(log""))
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792287908
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Yes, my concern is that decompress is called in every case, sometimes multiple times, which seems a bit redundant and makes the code somewhat hard to read. Below is a rough sketch (generated by ChatGPT).
```scala
private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
: (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
val userId = request.getUserContext.getUserId
val sessionId = request.getSessionId
val planAccessors: Seq[(proto.AnalyzePlanRequest => proto.Plan,
(proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest)] =
request.getAnalyzeCase match {
case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
Seq(
(_.getSchema.getPlan, (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
Seq(
(_.getExplain.getPlan, (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
Seq(
(_.getSameSemantics.getTargetPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build()),
(_.getSameSemantics.getOtherPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
)
// other single-plan cases...
case _ => Seq()
}
// Decompress all plans in sequence
var req = request
var sizes: Seq[Option[Long]] = Seq.empty
planAccessors.foreach { case (getter, setter) =>
val (newReq, size) = decompressPlanGeneric(req, userId, sessionId, getter, setter)
req = newReq
sizes = sizes :+ size
}
val targetSize = sizes.headOption.flatten
val otherSize = if (sizes.length > 1) sizes(1) else None
(req, targetSize, otherSize)
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on PR #54113:
URL: https://github.com/apache/spark/pull/54113#issuecomment-3840342690
Hi @hvanhovell , we're moving the decompression logic from AnalyzePlanHandler and ExecutePlanHandler into an interceptor. Could you take a look at this PR? Thanks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792287908
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Yes, the pattern matching of request.getAnalyzeCase cannot be avoided. But my concern is that decompress is called in every case, sometimes multiple times, which seems a bit redundant and makes the code somewhat hard to read. Below is a rough sketch (generated by ChatGPT).
```scala
private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
: (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
val userId = request.getUserContext.getUserId
val sessionId = request.getSessionId
val planAccessors: Seq[(proto.AnalyzePlanRequest => proto.Plan,
(proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest)] =
request.getAnalyzeCase match {
case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
Seq(
(_.getSchema.getPlan, (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
Seq(
(_.getExplain.getPlan, (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
)
case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
Seq(
(_.getSameSemantics.getTargetPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build()),
(_.getSameSemantics.getOtherPlan, (r, p) => r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
)
// other single-plan cases...
case _ => Seq()
}
// Decompress all plans in sequence
var req = request
var sizes: Seq[Option[Long]] = Seq.empty
planAccessors.foreach { case (getter, setter) =>
val (newReq, size) = decompressPlanGeneric(req, userId, sessionId, getter, setter)
req = newReq
sizes = sizes :+ size
}
val targetSize = sizes.headOption.flatten
val otherSize = if (sizes.length > 1) sizes(1) else None
(req, targetSize, otherSize)
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792393799
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
I am still a bit confused about the naming pattern (the first one should be named target-compressed-size if we follow it strictly). Still, I feel that storing the sizes in an array (if possible) could simplify the below logging logic:
```scala
// Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
val compressedSize = RequestDecompressionContext.getCompressedSize
val otherCompressedSize = RequestDecompressionContext.getOtherCompressedSize
if (compressedSize.isDefined || otherCompressedSize.isDefined) {
logDebug(
log"AnalyzePlan request received with compressed plan: " +
log"compressedSize=${MDC(BYTE_SIZE, compressedSize.getOrElse(0L))} bytes" +
otherCompressedSize
.map { size =>
log", otherCompressedSize=${MDC(BYTE_SIZE, size)} bytes"
}
.getOrElse(log""))
}
```
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
I am still a bit confused about the naming pattern (the first one should be named target-compressed-size if we follow it strictly). I still feel that storing the sizes in an array (if possible) could simplify the below logging logic:
```scala
// Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
val compressedSize = RequestDecompressionContext.getCompressedSize
val otherCompressedSize = RequestDecompressionContext.getOtherCompressedSize
if (compressedSize.isDefined || otherCompressedSize.isDefined) {
logDebug(
log"AnalyzePlan request received with compressed plan: " +
log"compressedSize=${MDC(BYTE_SIZE, compressedSize.getOrElse(0L))} bytes" +
otherCompressedSize
.map { size =>
log", otherCompressedSize=${MDC(BYTE_SIZE, size)} bytes"
}
.getOrElse(log""))
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2792393799
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
I am still a bit confused about the naming pattern (the first one should be named target-compressed-size if we follow it strictly). Still, I feel that storing the sizes in an array could simplify the below logging logic:
```scala
// Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
val compressedSize = RequestDecompressionContext.getCompressedSize
val otherCompressedSize = RequestDecompressionContext.getOtherCompressedSize
if (compressedSize.isDefined || otherCompressedSize.isDefined) {
logDebug(
log"AnalyzePlan request received with compressed plan: " +
log"compressedSize=${MDC(BYTE_SIZE, compressedSize.getOrElse(0L))} bytes" +
otherCompressedSize
.map { size =>
log", otherCompressedSize=${MDC(BYTE_SIZE, size)} bytes"
}
.getOrElse(log""))
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2794946393
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Hmm, another benefit I can foresee is that if other developers need to add getters/setters for a new AnalyzeCase, it will create less cognitive load when updating the code if we do the refactoring. Otherwise, they would still need to call decompress for req and finalReq (for the second plan, if needed), and pass the size/other size (for the second plan, if needed), instead of just figuring out how to implement the getter/setter.
In short, I am suggesting that we follow the DRY (Don't Repeat Yourself) principle.
@hvanhovell @vicennial WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2794641030
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
Yes, makes sense, it will be more clean. Updated to use a single key `COMPRESSED_SIZES_KEY` with `Seq[Option[Long]]`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "heyihong (via GitHub)" <gi...@apache.org>.
heyihong commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2794946393
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
Review Comment:
Hmm, another benefit I can foresee is that if other developers need to add getters/setters for a new AnalyzeCase, it will create less cognitive load when updating the code if we do the refactoring. Otherwise, they would still need to call decompress for reqWithTarget and finalReq (for the second plan), and pass the size/other size (for the second plan), instead of just figuring out how to write the getter/setter.
@hvanhovell @vicennial WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2787688891
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,302 @@
+/*
+ * 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.connect.service
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create a listener that will intercept and decompress the message
+ val listener = next.startCall(call, headers)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Option[Long], Option[Long])): Unit = {
+ val (decompressedReq, compressedSize, otherCompressedSize) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed size(s) in context if present
+ var contextToUse = Context.current()
+ compressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.COMPRESSED_SIZE_KEY, size)
+ }
+ otherCompressedSize.foreach { size =>
+ contextToUse =
+ contextToUse.withValue(RequestDecompressionContext.OTHER_COMPRESSED_SIZE_KEY, size)
+ }
+
+ // Run the rest of the call chain with the context
+ val prev = contextToUse.attach()
+ try {
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ } finally {
+ contextToUse.detach(prev)
+ }
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(request: proto.ExecutePlanRequest)
+ : (proto.ExecutePlanRequest, Option[Long], Option[Long]) = {
+ if (!request.hasPlan) {
+ return (request, None, None)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, size, None)
+ }
+
+ private def decompressAnalyzePlanRequest(request: proto.AnalyzePlanRequest)
+ : (proto.AnalyzePlanRequest, Option[Long], Option[Long]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, size, None)
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, targetSize, otherSize)
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, None, None)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, None, None)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, None, None)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to ExecuteHolder. Uses gRPC
+ * Context to properly propagate values across async boundaries.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing the compressed size. This is set by RequestDecompressionInterceptor
+ * when a compressed request is encountered, and read by handlers for metrics.
+ */
+ val COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("compressed-size")
+
+ /**
+ * Context key for storing the other compressed size. This is used for AnalyzePlanRequest
+ * SameSemantics which has two plans: target_plan (size stored in COMPRESSED_SIZE_KEY) and
+ * other_plan (size stored in this key).
+ */
+ val OTHER_COMPRESSED_SIZE_KEY: Context.Key[Long] = Context.key("other-compressed-size")
Review Comment:
Actually the name comes from the Spark Connect protocol: the two values correspond to target_plan and other_plan: https://github.com/apache/spark/blob/ebd5b007fcf203eadcf8b037ab2b99577490f869/sql/connect/common/src/main/protobuf/spark/connect/base.proto#L190-L196
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2878447390
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,316 @@
+/*
+ * 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.connect.service
+
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Contexts, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create an AtomicReference to hold compressed sizes
+ val compressedSizesRef = new AtomicReference[Seq[Option[Long]]](Seq.empty)
+
+ // Create context with the AtomicReference and start the call within that context
+ val ctx = Context
+ .current()
+ .withValue(RequestDecompressionContext.COMPRESSED_SIZES_KEY, compressedSizesRef)
+ val listener = Contexts.interceptCall(ctx, call, headers, next)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ compressedSizesRef,
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ compressedSizesRef,
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ compressedSizesRef: AtomicReference[Seq[Option[Long]]],
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Seq[Option[Long]])): Unit = {
+ val (decompressedReq, compressedSizes) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed sizes in the AtomicReference
+ compressedSizesRef.set(compressedSizes)
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(
+ request: proto.ExecutePlanRequest): (proto.ExecutePlanRequest, Seq[Option[Long]]) = {
+ if (!request.hasPlan) {
+ return (request, Seq.empty)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, Seq(size))
+ }
+
+ private def decompressAnalyzePlanRequest(
+ request: proto.AnalyzePlanRequest): (proto.AnalyzePlanRequest, Seq[Option[Long]]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, Seq(targetSize, otherSize))
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, Seq.empty)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, Seq.empty)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, Seq.empty)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
+ request: Req,
+ userId: String,
+ sessionId: String,
+ getPlan: Req => proto.Plan,
+ rebuild: (Req, proto.Plan) => Req): (Req, Option[Long]) = {
+
+ val plan = getPlan(request)
+
+ if (plan.getOpTypeCase != proto.Plan.OpTypeCase.COMPRESSED_OPERATION) {
+ return (request, None)
+ }
+ val compressedSize = plan.getCompressedOperation.getData.size().toLong
+ val requestType = request.getClass.getSimpleName
+ logInfo(
+ log"Received compressed plan in ${MDC(CLASS_NAME, requestType)} " +
+ log"(size=${MDC(BYTE_SIZE, compressedSize)} bytes): " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedPlan = PlanCompressionUtils.decompressPlan(plan)
+
+ logInfo(
+ log"Plan decompression completed " +
+ log"(compressed=${MDC(BYTE_SIZE, compressedSize)} bytes -> " +
+ log"decompressed=${MDC(BYTE_SIZE, decompressedPlan.getSerializedSize.toLong)} bytes, " +
+ log"userId=${MDC(USER_ID, userId)}, sessionId=${MDC(SESSION_ID, sessionId)}")
+
+ val decompressedRequest = rebuild(request, decompressedPlan)
+ (decompressedRequest, Some(compressedSize))
+ }
+}
+
+/**
+ * Context holder for passing decompression metrics from interceptor to handlers. Uses gRPC
+ * Context to properly propagate values.
+ */
+object RequestDecompressionContext {
+
+ /**
+ * Context key for storing compressed sizes. This is set by RequestDecompressionInterceptor when
+ * compressed requests are encountered, and read by handlers for metrics.
+ *
+ * The sequence contains Option[Long] entries corresponding to each plan in the request:
+ * - For ExecutePlan and single-plan Analyze requests: Seq(Some(size)) if compressed, or
+ * Seq.empty if not
+ * - For AnalyzePlanRequest SameSemantics with two plans: Seq(target_size, other_size) where
+ * each is Some(size) if that plan was compressed, None if not
+ *
+ * The sequence length matches the number of plans, with explicit None for uncompressed plans.
+ */
+ val COMPRESSED_SIZES_KEY: Context.Key[AtomicReference[Seq[Option[Long]]]] =
Review Comment:
You’re right that gRPC Contexts are thread-local, but the Context is supposed to be set in `interceptCall`, not in `SimpleForwardingServerCallListener.onMessage` when we have a request. If I set the context in `onMessage`, I find that the context I set always gets overridden back to the original one and cannot be propagated to the handler, so I can only set it as an `AtomicReference` in `interceptCall` and then update the value in `onMessage`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2878455915
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/RequestDecompressionInterceptor.scala:
##########
@@ -0,0 +1,316 @@
+/*
+ * 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.connect.service
+
+import java.util.concurrent.atomic.AtomicReference
+
+import scala.util.control.NonFatal
+
+import io.grpc.{Context, Contexts, Metadata, ServerCall, ServerCallHandler, ServerInterceptor}
+import io.grpc.ForwardingServerCallListener.SimpleForwardingServerCallListener
+
+import org.apache.spark.connect.proto
+import org.apache.spark.internal.Logging
+import org.apache.spark.internal.LogKeys.{BYTE_SIZE, CLASS_NAME, SESSION_ID, USER_ID}
+import org.apache.spark.sql.connect.utils.{ErrorUtils, PlanCompressionUtils}
+
+/**
+ * Interceptor that decompresses compressed requests before they reach downstream handlers.
+ *
+ * This interceptor currently handles:
+ * - ExecutePlanRequest with compressed plans
+ * - AnalyzePlanRequest with compressed plans (Schema, Explain, TreeString, IsLocal,
+ * IsStreaming, InputFiles, SemanticHash, and SameSemantics analysis types)
+ *
+ * Compressed plan size metrics are tracked in gRPC Context for use by handlers.
+ */
+class RequestDecompressionInterceptor extends ServerInterceptor with Logging {
+
+ override def interceptCall[ReqT, RespT](
+ call: ServerCall[ReqT, RespT],
+ headers: Metadata,
+ next: ServerCallHandler[ReqT, RespT]): ServerCall.Listener[ReqT] = {
+
+ // Create an AtomicReference to hold compressed sizes
+ val compressedSizesRef = new AtomicReference[Seq[Option[Long]]](Seq.empty)
+
+ // Create context with the AtomicReference and start the call within that context
+ val ctx = Context
+ .current()
+ .withValue(RequestDecompressionContext.COMPRESSED_SIZES_KEY, compressedSizesRef)
+ val listener = Contexts.interceptCall(ctx, call, headers, next)
+
+ new SimpleForwardingServerCallListener[ReqT](listener) {
+ override def onMessage(message: ReqT): Unit = {
+ message match {
+ case req: proto.ExecutePlanRequest =>
+ handleRequestWithDecompression(
+ compressedSizesRef,
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressExecutePlanRequest(req))
+
+ case req: proto.AnalyzePlanRequest =>
+ handleRequestWithDecompression(
+ compressedSizesRef,
+ req.getUserContext.getUserId,
+ req.getSessionId,
+ () => decompressAnalyzePlanRequest(req))
+
+ case other =>
+ // Forward all other message types as-is (no decompression or error handling needed)
+ super.onMessage(other)
+ }
+ }
+
+ private def handleRequestWithDecompression[T](
+ compressedSizesRef: AtomicReference[Seq[Option[Long]]],
+ userId: String,
+ sessionId: String,
+ decompressRequest: () => (T, Seq[Option[Long]])): Unit = {
+ val (decompressedReq, compressedSizes) =
+ try {
+ decompressRequest()
+ } catch {
+ case NonFatal(e) =>
+ // Handle decompression errors
+ logError(
+ log"Plan decompression failed: " +
+ log"userId=${MDC(USER_ID, userId)}, " +
+ log"sessionId=${MDC(SESSION_ID, sessionId)}",
+ e)
+ ErrorUtils.handleError("planDecompression", call, userId, sessionId)(e)
+ return
+ }
+
+ // Set compressed sizes in the AtomicReference
+ compressedSizesRef.set(compressedSizes)
+ super.onMessage(decompressedReq.asInstanceOf[ReqT])
+ }
+ }
+ }
+
+ private def decompressExecutePlanRequest(
+ request: proto.ExecutePlanRequest): (proto.ExecutePlanRequest, Seq[Option[Long]]) = {
+ if (!request.hasPlan) {
+ return (request, Seq.empty)
+ }
+ val (decompressedReq, size) = decompressPlanGeneric(
+ request,
+ request.getUserContext.getUserId,
+ request.getSessionId,
+ (r: proto.ExecutePlanRequest) => r.getPlan,
+ (req: proto.ExecutePlanRequest, plan: proto.Plan) => req.toBuilder.setPlan(plan).build())
+
+ (decompressedReq, Seq(size))
+ }
+
+ private def decompressAnalyzePlanRequest(
+ request: proto.AnalyzePlanRequest): (proto.AnalyzePlanRequest, Seq[Option[Long]]) = {
+ val userId = request.getUserContext.getUserId
+ val sessionId = request.getSessionId
+
+ // Helper to decompress a single-plan analysis type
+ def decompress(
+ req: proto.AnalyzePlanRequest,
+ getPlan: proto.AnalyzePlanRequest => proto.Plan,
+ rebuild: (proto.AnalyzePlanRequest, proto.Plan) => proto.AnalyzePlanRequest) = {
+ decompressPlanGeneric(req, userId, sessionId, getPlan, rebuild)
+ }
+
+ // NOTE: All AnalyzePlanRequest cases are explicitly listed here.
+ // The default case throws an exception to catch new cases at runtime and fail CI tests.
+ request.getAnalyzeCase match {
+ // Cases with Plan fields - decompress if compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.SCHEMA =>
+ val (req, size) = decompress(
+ request,
+ _.getSchema.getPlan,
+ (r, p) => r.toBuilder.setSchema(r.getSchema.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.EXPLAIN =>
+ val (req, size) = decompress(
+ request,
+ _.getExplain.getPlan,
+ (r, p) => r.toBuilder.setExplain(r.getExplain.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.TREE_STRING =>
+ val (req, size) = decompress(
+ request,
+ _.getTreeString.getPlan,
+ (r, p) => r.toBuilder.setTreeString(r.getTreeString.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_LOCAL =>
+ val (req, size) = decompress(
+ request,
+ _.getIsLocal.getPlan,
+ (r, p) => r.toBuilder.setIsLocal(r.getIsLocal.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.IS_STREAMING =>
+ val (req, size) = decompress(
+ request,
+ _.getIsStreaming.getPlan,
+ (r, p) => r.toBuilder.setIsStreaming(r.getIsStreaming.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.INPUT_FILES =>
+ val (req, size) = decompress(
+ request,
+ _.getInputFiles.getPlan,
+ (r, p) => r.toBuilder.setInputFiles(r.getInputFiles.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SEMANTIC_HASH =>
+ val (req, size) = decompress(
+ request,
+ _.getSemanticHash.getPlan,
+ (r, p) => r.toBuilder.setSemanticHash(r.getSemanticHash.toBuilder.setPlan(p)).build())
+ (req, Seq(size))
+
+ case proto.AnalyzePlanRequest.AnalyzeCase.SAME_SEMANTICS =>
+ // Special case: has two Plan fields (target_plan and other_plan)
+ val (reqWithTarget, targetSize) = decompress(
+ request,
+ _.getSameSemantics.getTargetPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setTargetPlan(p)).build())
+
+ val (finalReq, otherSize) = decompress(
+ reqWithTarget,
+ _.getSameSemantics.getOtherPlan,
+ (r, p) =>
+ r.toBuilder.setSameSemantics(r.getSameSemantics.toBuilder.setOtherPlan(p)).build())
+
+ (finalReq, Seq(targetSize, otherSize))
+
+ // Cases with Relation fields - currently not compressed
+ case proto.AnalyzePlanRequest.AnalyzeCase.PERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.UNPERSIST |
+ proto.AnalyzePlanRequest.AnalyzeCase.GET_STORAGE_LEVEL =>
+ (request, Seq.empty)
+
+ // Cases with no Plan or Relation fields - safe to pass through
+ case proto.AnalyzePlanRequest.AnalyzeCase.SPARK_VERSION |
+ proto.AnalyzePlanRequest.AnalyzeCase.DDL_PARSE |
+ proto.AnalyzePlanRequest.AnalyzeCase.JSON_TO_DDL =>
+ (request, Seq.empty)
+
+ // No analysis case set - safe to pass through, will be handled in handler
+ case proto.AnalyzePlanRequest.AnalyzeCase.ANALYZE_NOT_SET =>
+ (request, Seq.empty)
+
+ case _ =>
+ // Unhandled case - fail to catch new cases during testing
+ throw new UnsupportedOperationException(
+ s"Unhandled AnalyzePlanRequest case: ${request.getAnalyzeCase}. " +
+ s"RequestDecompressionInterceptor must be updated to handle this case explicitly. " +
+ s"If the case contains Plan fields, add decompression logic. " +
+ s"Otherwise, add it to the safe passthrough cases.")
+ }
+ }
+
+ private def decompressPlanGeneric[Req](
Review Comment:
Makes sense! Replaced the generic `decompressPlanGeneric[Req]` (which took getPlan/rebuild callbacks and operated on requests) with a simpler `decompressPlan` that only takes a proto.Plan and decompress.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on code in PR #54113:
URL: https://github.com/apache/spark/pull/54113#discussion_r2878466512
##########
sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectAnalyzeHandler.scala:
##########
@@ -69,9 +70,23 @@ private[connect] class SparkConnectAnalyzeHandler(
throw new AnalysisException("ATTEMPT_ANALYSIS_IN_PIPELINE_QUERY_FUNCTION", Map())
}
+ // Log compressed sizes from gRPC Context (set by RequestDecompressionInterceptor)
+ val compressedSize = RequestDecompressionContext.getCompressedSize
+ val otherCompressedSize = RequestDecompressionContext.getOtherCompressedSize
+ if (compressedSize.isDefined || otherCompressedSize.isDefined) {
+ logDebug(
+ log"AnalyzePlan request received with compressed plan: " +
+ log"compressedSize=${MDC(BYTE_SIZE, compressedSize.getOrElse(0L))} bytes" +
+ otherCompressedSize
+ .map { size =>
+ log", otherCompressedSize=${MDC(BYTE_SIZE, size)} bytes"
+ }
+ .getOrElse(log""))
+ }
+
def transformRelation(rel: proto.Relation) = planner.transformRelation(rel, cachePlan = true)
def transformRelationPlan(plan: proto.Plan) = {
Review Comment:
Yes, but I would like to keep it because if we want to have plan-level operations in the future we will need this tiny helper back. Changed this to one-line `def transformRelationPlan(plan: proto.Plan) = transformRelation(plan.getRoot)` to be simple.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-55338][CONNECT] Centralize Spark Connect request decompression logic in gRPC interceptor [spark]
Posted by "xi-db (via GitHub)" <gi...@apache.org>.
xi-db commented on PR #54113:
URL: https://github.com/apache/spark/pull/54113#issuecomment-3853942813
cc @vicennial Here's a followup PR for the plan compression. We're moving the decompression in different handlers into a centralized gRPC interceptor.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org