Spark Tasks Failing Intermittently With Connection Reset Error - sql-server

Question

I am using the spark dataframe reader to pull data from an SQL server database, do some minor changes like column renaming, data type casting and saving the dataframe to S3 using delta lake format. The job is triggered from Airflow using LivyOperator.
The code for the dataframe reader is as follows:
val table_source = spark.read
.format("com.microsoft.sqlserver.jdbc.spark")
.option("Driver", "com.microsoft.sqlserver.jdbc.SQLServerDriver")
.option("dbtable", select_query)
.option("inferschema", "true")
.option("url", ds_url)
.option("user", ds_user)
.option("password", ds_pass)
.option("numPartitions", num_partitions)
.option("partitionColumn", "RNO")
.option("lowerBound", 0)
.option("upperBound", rowCount).load()
Then I create a temporary view on top of this data and add a few more standard columns like client_id, timestamp etc and return a dataframe. Then the dataframe is saved as parquet files using 'delta' format.
table_source.createOrReplaceTempView("table_name")
val table_mapping = spark.sql(mapping_query)
table_mapping
.write.format("delta")
.mode("append")
.save(path)
So, now the issue is that for tables that have around 50k rows or more the job always gets stuck at this 'save' stage. Say I provided numPartitions=8, in the spark UI I can see that 8 tasks are generated for this stage by spark. Some of these tasks finish executing successfully within a few minutes. The remaining tasks get stuck for more than 2 hours and fail with a 'connection reset' error caused by SQLServerException. The tasks are then retried by spark, some of these tasks complete right away and again some get stuck for two more hours and so on until the stage is completed eventually.
Note: There is no limit on the maximum concurrent connections on the source server.
Finally job takes around 2+, 4+ or 6+ hours to complete. Using spark speculation helped bring the job completion time to 1 hour, but it is still too high a time for the volume of data we are dealing with. For comparison we tested fetching the data in the same environment as the EMR cluster using SSIS. It took just 20 mins to complete.
When the tasks are stuck we observed the following thread lock in the executor thread dump.
Thread ID: ##
Thread Name: Executor task launch worker for task 1.0 in stage 17.0 (TID 17)
Thread State: RUNNABLE
Thread Locks: Lock(java.util.concurrent.ThreadPoolExecutor$Worker#881566968}),
Monitor(com.microsoft.sqlserver.jdbc.TDSReader#1700342603})
When I expand this thread, I see the following trace.
> java.net.SocketInputStream.socketRead0(Native Method)
> java.net.SocketInputStream.socketRead(SocketInputStream.java:116)
> java.net.SocketInputStream.read(SocketInputStream.java:171)
> java.net.SocketInputStream.read(SocketInputStream.java:141)
> com.microsoft.sqlserver.jdbc.TDSChannel.read(IOBuffer.java:1819)
> com.microsoft.sqlserver.jdbc.TDSReader.readPacket(IOBuffer.java:5461)
> => holding Monitor(com.microsoft.sqlserver.jdbc.TDSReader#648027762}) com.microsoft.sqlserver.jdbc.TDSReader.nextPacket(IOBuffer.java:5371)
> com.microsoft.sqlserver.jdbc.TDSReader.ensurePayload(IOBuffer.java:5347)
> com.microsoft.sqlserver.jdbc.TDSReader.readBytes(IOBuffer.java:5640)
> com.microsoft.sqlserver.jdbc.TDSReader.readWrappedBytes(IOBuffer.java:5662)
> com.microsoft.sqlserver.jdbc.TDSReader.readInt(IOBuffer.java:5603)
> com.microsoft.sqlserver.jdbc.TDSReader.readUnsignedInt(IOBuffer.java:5620)
> com.microsoft.sqlserver.jdbc.PLPInputStream.readBytesInternal(PLPInputStream.java:313)
> com.microsoft.sqlserver.jdbc.PLPInputStream.getBytes(PLPInputStream.java:129)
> com.microsoft.sqlserver.jdbc.DDC.convertStreamToObject(DDC.java:438)
> com.microsoft.sqlserver.jdbc.ServerDTVImpl.getValue(dtv.java:2965)
> com.microsoft.sqlserver.jdbc.DTV.getValue(dtv.java:206)
> com.microsoft.sqlserver.jdbc.Column.getValue(Column.java:130)
> com.microsoft.sqlserver.jdbc.SQLServerResultSet.getValue(SQLServerResultSet.java:2087)
> com.microsoft.sqlserver.jdbc.SQLServerResultSet.getValue(SQLServerResultSet.java:2072)
> com.microsoft.sqlserver.jdbc.SQLServerResultSet.getString(SQLServerResultSet.java:2413)
> org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils$.$anonfun$makeGetter$12(JdbcUtils.scala:444)
> org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils$.$anonfun$makeGetter$12$adapted(JdbcUtils.scala:442)
> org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils$$$Lambda$1086/1697796400.apply(Unknown
> Source)
> org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils$$anon$1.getNext(JdbcUtils.scala:352)
> org.apache.spark.sql.execution.datasources.jdbc.JdbcUtils$$anon$1.getNext(JdbcUtils.scala:334)
> org.apache.spark.util.NextIterator.hasNext(NextIterator.scala:73)
> org.apache.spark.InterruptibleIterator.hasNext(InterruptibleIterator.scala:37)
> org.apache.spark.util.CompletionIterator.hasNext(CompletionIterator.scala:31)
> org.apache.spark.sql.catalyst.expressions.GeneratedClass$GeneratedIteratorForCodegenStage1.processNext(Unknown
> Source)
> org.apache.spark.sql.execution.BufferedRowIterator.hasNext(BufferedRowIterator.java:35)
> org.apache.spark.sql.execution.WholeStageCodegenExec$$anon$1.hasNext(WholeStageCodegenExec.scala:832)
> org.apache.spark.sql.execution.datasources.FileFormatWriter$.$anonfun$executeTask$1(FileFormatWriter.scala:277)
> org.apache.spark.sql.execution.datasources.FileFormatWriter$$$Lambda$1243/1672950527.apply(Unknown
> Source)
> org.apache.spark.util.Utils$.tryWithSafeFinallyAndFailureCallbacks(Utils.scala:1473)
> org.apache.spark.sql.execution.datasources.FileFormatWriter$.executeTask(FileFormatWriter.scala:286)
> org.apache.spark.sql.execution.datasources.FileFormatWriter$.$anonfun$write$15(FileFormatWriter.scala:210)
> org.apache.spark.sql.execution.datasources.FileFormatWriter$$$Lambda$1085/1036621064.apply(Unknown
> Source)
> org.apache.spark.scheduler.ResultTask.runTask(ResultTask.scala:90)
> org.apache.spark.scheduler.Task.run(Task.scala:131)
> org.apache.spark.executor.Executor$TaskRunner.$anonfun$run$3(Executor.scala:497)
> org.apache.spark.executor.Executor$TaskRunner$$Lambda$465/565856309.apply(Unknown
> Source)
> org.apache.spark.util.Utils$.tryWithSafeFinally(Utils.scala:1439)
> org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:500)
> java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1149)
> java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:624)
> java.lang.Thread.run(Thread.java:750)
>
>
> Version Information
> EMR: 6.40
> Spark: 3.1.2
> mssql-jdbc: 9.2.1.jre8
> spark-mssql-connector_2.12: 1.2.0
> delta-core: 1.0.0
We have tried setting QueryTimeout on the spark jdbc reader to no avail. Tried increasing executor and driver cores/memory and also using dynamic allocation but ended up with the same results. Also tried removing partitions, same issue. We have been at this for weeks, I would highly appreciate any pointers on solving this problem.