Skip to content

[SPARK-49673][CONNECT] Increase CONNECT_GRPC_ARROW_MAX_BATCH_SIZE to 0.7 * CONNECT_GRPC_MAX_MESSAGE_SIZE #48122

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 8 commits into from
Closed

[SPARK-49673][CONNECT] Increase CONNECT_GRPC_ARROW_MAX_BATCH_SIZE to 0.7 * CONNECT_GRPC_MAX_MESSAGE_SIZE #48122

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
a68f280
increase maxBatchSize to 0.7 * grpc message size limit for sqlCommand…
dillitz Sep 16, 2024
18a330f
remove unused import
dillitz Sep 16, 2024
c04795c
Revert "remove unused import"
dillitz Sep 17, 2024
cf006e8
Revert "increase maxBatchSize to 0.7 * grpc message size limit for sq…
dillitz Sep 17, 2024
01eddd3
change CONNECT_GRPC_ARROW_MAX_BATCH_SIZE to CONNECT_GRPC_MAX_MESSAGE_…
dillitz Sep 17, 2024
4eb6711
add E2E test
dillitz Sep 17, 2024
ca10ef9
adjust client's maxMessageSize
dillitz Sep 17, 2024
b0bfb27
linter
dillitz Sep 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 21 additions & 2 deletions connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/ClientE2ETestSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import java.util.Properties

import scala.collection.mutable
import scala.concurrent.{ExecutionContext, Future}
import scala.concurrent.duration.DurationInt
import scala.concurrent.duration.{DurationInt, FiniteDuration}
import scala.jdk.CollectionConverters._

import org.apache.commons.io.FileUtils
Expand All @@ -37,7 +37,7 @@ import org.apache.spark.sql.catalyst.analysis.{NamespaceAlreadyExistsException,
import org.apache.spark.sql.catalyst.encoders.AgnosticEncoders.StringEncoder
import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
import org.apache.spark.sql.catalyst.parser.ParseException
import org.apache.spark.sql.connect.client.{SparkConnectClient, SparkResult}
import org.apache.spark.sql.connect.client.{RetryPolicy, SparkConnectClient, SparkResult}
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SqlApiConf
import org.apache.spark.sql.test.{ConnectFunSuite, IntegrationTestUtils, RemoteSparkSession, SQLHelper}
Expand Down Expand Up @@ -1566,6 +1566,25 @@ class ClientE2ETestSuite
val result = df.select(trim(col("col"), " ").as("trimmed_col")).collect()
assert(result sameElements Array(Row("a"), Row("b"), Row("c")))
}

test("SPARK-49673: new batch size, multiple batches") {
val maxBatchSize = spark.conf.get("spark.connect.grpc.arrow.maxBatchSize").dropRight(1).toInt
// Adjust client grpcMaxMessageSize to maxBatchSize (10MiB; set in RemoteSparkSession config)
val sparkWithLowerMaxMessageSize = SparkSession
.builder()
.client(
SparkConnectClient
.builder()
.userId("test")
.port(port)
.grpcMaxMessageSize(maxBatchSize)
.retryPolicy(RetryPolicy
.defaultPolicy()
.copy(maxRetries = Some(10), maxBackoff = Some(FiniteDuration(30, "s"))))
.build())
.create()
assert(sparkWithLowerMaxMessageSize.range(maxBatchSize).collect().length == maxBatchSize)
}
}

private[sql] case class ClassData(a: String, b: Int)
Expand Down
2 changes: 2 additions & 0 deletions ...ctor/connect/client/jvm/src/test/scala/org/apache/spark/sql/test/RemoteSparkSession.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,8 @@ object SparkConnectServerUtils {
// to make the tests exercise reattach.
"spark.connect.execute.reattachable.senderMaxStreamDuration=1s",
"spark.connect.execute.reattachable.senderMaxStreamSize=123",
// Testing SPARK-49673, setting maxBatchSize to 10MiB
s"spark.connect.grpc.arrow.maxBatchSize=${10 * 1024 * 1024}",
// Disable UI
"spark.ui.enabled=false")
Seq("--jars", catalystTestJar) ++ confs.flatMap(v => "--conf" :: v :: Nil)
Expand Down
2 changes: 1 addition & 1 deletion sql/connect/server/src/main/scala/org/apache/spark/sql/connect/config/Connect.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ object Connect {
"conservatively use 70% of it because the size is not accurate but estimated.")
.version("3.4.0")
.bytesConf(ByteUnit.BYTE)
.createWithDefault(4 * 1024 * 1024)
.createWithDefault(ConnectCommon.CONNECT_GRPC_MAX_MESSAGE_SIZE)

val CONNECT_GRPC_MAX_INBOUND_MESSAGE_SIZE =
buildStaticConf("spark.connect.grpc.maxInboundMessageSize")
Expand Down