dongjoon-hyun opened a new pull request, #59945:
URL: https://github.com/apache/spark/pull/59945
### What changes were proposed in this pull request?
This PR fixes an Arrow memory leak in Spark Connect when the conversion of
results to Arrow batches on the driver fails or is interrupted.
`SparkConnectPlanExecution.processAsArrowBatches` converts the results of
`LocalTableScanExec`, `CollectLimitExec` and `CollectTailExec` on the driver,
and `SparkConnectPlanner.handleSqlCommand` converts the results of commands and
SQL scripts on the driver. Both use
`ArrowConverters.toBatchWithSchemaIterator`, whose iterator allocates a child
allocator of `ArrowUtils.rootAllocator` and the vectors of a
`VectorSchemaRoot`. In a task, a task completion listener closes the iterator.
On the driver, there is no `TaskContext`, so the iterator is closed only when
`hasNext` reaches the end of the rows. If `next()` or sending a batch throws,
nothing closes the iterator, and its memory stays in `ArrowUtils.rootAllocator`
until the driver stops.
This PR:
- makes `SparkConnectPlanExecution.rowToArrowConverter` return a
`CloseableIterator` and closes it with `Utils.tryWithSafeFinally` when the rows
are converted on the driver. Tasks still rely on the task completion listener.
- closes the iterator in `SparkConnectPlanner.handleSqlCommand` in the same
way.
- closes the `ArrowRecordBatch` in a `finally` block in
`ArrowBatchWithSchemaIterator.next()`, as `ArrowBatchIterator.next()` already
does. Otherwise, when `MessageSerializer.serialize` throws, the buffers that
the batch retains stay allocated, and closing the iterator fails with
`IllegalStateException: Memory was leaked by query` instead of releasing them.
`Utils.tryWithSafeFinally` keeps the original exception and adds a failure
of `close()` to it as a suppressed exception. Closing the iterator again after
`hasNext` has closed it is a no-op.
Alternatives that were not taken:
- Make the iterator close itself when `next()` fails: it does not cover a
failure while sending a batch.
- Make `ArrowBatchIterator` extend `CloseableIterator` so that `map` keeps
`close()`: it changes the type hierarchy of a `sql/core` class that tasks also
use, only for the driver-side paths of Spark Connect.
### Why are the changes needed?
To fix a memory leak in the Spark Connect server. The most common trigger is
an interrupt. `ExecuteThreadRunner.interrupt()` interrupts the execution
thread, for example when the PySpark client sends `Interrupt` on
`KeyboardInterrupt` (Ctrl-C), on `interruptAll`, `interruptTag` and
`interruptOperation`, on `ReleaseExecute` with `release_all` (e.g. when a
client closes the result iterator early), and when the session is closed.
`ArrowBatchWithSchemaIterator.next()` writes to a `WriteChannel` over
`Channels.newChannel(ByteArrayOutputStream)`, which is an
`AbstractInterruptibleChannel`, so its next write throws
`ClosedByInterruptException`. The execution thread spends most of its time
filling a batch, so the exception usually comes from
`MessageSerializer.serialize(writeChannel, batch)` and also leaves the
`ArrowRecordBatch` open.
An error from `ArrowWriter` leaks in the same way. For example, the
following query fails with `DATETIME_OVERFLOW` because the microseconds of the
interval do not fit in the int64 nanoseconds of Arrow, and it leaves 65536
bytes in `ArrowUtils.rootAllocator` of the driver.
```sql
SELECT * FROM VALUES (make_interval(0, 0, 0, 0, 2600000, 0, 0)) AS t(i)
```
In a local test with the default
`spark.sql.execution.arrow.maxRecordsPerBatch`, each interrupt of `SELECT id,
repeat('x', 200) AS s FROM range(0, 400000, 1, 1) LIMIT 300000` leaked
2,293,760 bytes. A long-running Spark Connect server accumulates these leaks in
direct memory. SPARK-45814 (#43691) fixed the same kind of leak in
`ArrowConverters.createEmptyArrowBatch`.
By code reading, the conversion in `handleSqlCommand` has existed since
Apache Spark 3.4.0 (SPARK-41725), the `LocalTableScanExec` path since Apache
Spark 3.5.0 (SPARK-43153), and the `CollectLimitExec` and `CollectTailExec`
paths since Apache Spark 4.2.0 (SPARK-55887).
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Pass the CIs with the newly added `SparkConnectPlanExecutionArrowLeakSuite`.
Each test checks that no `toArrowBatchWithSchemaIterator` allocator and no
memory remain in `ArrowUtils.rootAllocator`.
- `LocalTableScanExec`, `CollectLimitExec` and `CollectTailExec`: an
interrupt after the first sent batch, a failure while sending the first batch,
and a complete conversion
- `LocalTableScanExec`: an interrupt while filling a batch, an `ArrowWriter`
error, and an empty relation
- `CollectLimitExec`: an `ArrowWriter` error
- `ArrowBatchWithSchemaIterator`: closing the iterator after a failure while
serializing the batch
- End-to-end through `SparkConnectService`: an `ArrowWriter` error in a
local relation and in a SQL script result (`handleSqlCommand`), and an
`Interrupt` RPC while a `CollectLimitExec` or `LocalRelation` result is being
converted
Without this fix, all of them fail except the four tests of the complete
conversions and the empty relation.
I also ran `ArrowConvertersSuite`, `SparkConnectServiceSuite` and
`SparkConnectPlannerSuite` locally, and all of them passed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 5.5
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
