FelixYBW opened a new pull request, #13419:
URL: https://github.com/apache/gluten/pull/13419
## What changes are proposed in this pull request?
Part of #13381 (2 of 2). Depends on #13415, whose commits this branch
includes. Only the last commit is new here.
#13415 stops operators from reserving Spark memory again for Arrow batches
they load from native batches, by loading into a separate untracked allocator
root. `ColumnarArrowEvalPythonExec` couldn't use it. It took Java Arrow batches
(`LoadArrowDataExec` below it), joined the cached input columns with the Python
results into one Java batch, and emitted that (`OffloadArrowDataExec` above
it). Arrow can't load buffers from two allocator roots into one
`VectorSchemaRoot`, so the loaded batches stayed in the task allocator, and
their memory was reserved twice.
`ColumnarArrowEvalPythonExec` now takes and produces Velox batches, like
`ColumnarPartialProjectExec`:
- **Input**: retains each input batch and loads only the columns the UDFs
read (`ColumnarBatches.select`), into
`ArrowBufferAllocators.contextUntrackedInstance()`. The runner closes them once
they're written to the Python worker. Before, every input column was loaded and
kept until the result came back.
- **Output**: offloads the Python results and merges them into the retained
input batch natively (`VeloxColumnarBatches.compose`), so the input columns
never leave Velox. Each input batch is released once its result is merged.
Before, all input batches of a partition were held until the partition ended.
- **No Arrow transitions around the UDF**: the plan goes from `Velox op →
LoadArrowData → ColumnarArrowEvalPython → OffloadArrowData →
ArrowColumnarToVeloxColumnar → Velox op` to `Velox op → ColumnarArrowEvalPython
→ Velox op`.
- **`ColumnarArrowPythonRunner`**:
- The writer serializes each input batch directly (`MessageSerializer`),
instead of loading it into a `VectorSchemaRoot` of the task allocator, which
fails across allocator roots.
- The reader moves each result batch's buffers into new vectors (a
transfer within the same allocator, no copy), so the batch owns them. It no
longer reuses the root's vectors, so `adjustRefCnt` is removed.
- **Docs**: the Arrow UDF example in `ArrowColumnarConversion.md` is
rewritten for the new flow. The PlantUML source is updated, but
`docs/image/VeloxColumnArrowUdfLifecycle.png` isn't regenerated yet.
## How was this patch tested?
- `ArrowEvalPythonExecSuite`: the chained-UDF test now checks that there's
no `LoadArrowDataExec` below and no `OffloadArrowDataExec` above the UDFs.
- Spark 3.5 (pyspark 3.5.5, pandas 2.2.3, pyarrow 20.0.0, as in CI), all
passing: `ColumnarBatchTest`, `ArrowEvalPythonExecSuite`,
`UDFPartialProjectSuite`, `VeloxTransitionSuite`.
- Spark 4.0 (ANSI off, as in CI), all passing: `ArrowEvalPythonExecSuite`
(all 5 tests, including the chained-UDF and parquet-scan tests that only run on
4.0+), `UDFPartialProjectSuite`, `VeloxTransitionSuite`.
## Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (claude-opus-5-5)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
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]
